Spaces:
Runtime error
Runtime error
File size: 21,322 Bytes
c54f817 8435838 c575e18 89b9813 8435838 89b9813 e62ac65 e8de28c a677593 5a39c97 a2d3df0 b916cdf 4994a33 09eba18 b916cdf 8435838 fc0af52 5754a1c a9d9852 5754a1c a9d9852 5754a1c a9d9852 5754a1c 929432c 5754a1c 929432c 5754a1c a9d9852 929432c a9d9852 929432c 5754a1c 984922b 5754a1c 984922b dc01bcc a9d9852 984922b a9d9852 fc0af52 b403864 a9d9852 5754a1c a9d9852 5754a1c 8435838 c575e18 89b9813 db037f3 89b9813 fc0af52 a9930a8 89b9813 fc0af52 5587923 fc0af52 89b9813 fc0af52 db037f3 fc0af52 89b9813 a677593 0b874ec 8a16fb0 0b874ec e8de28c 3208a74 50eea75 eac6a5b c94bfc7 50eea75 505d2c9 50eea75 3208a74 5a39c97 f1ee247 50eea75 f1ee247 a9d9852 f1ee247 a9d9852 50eea75 f1ee247 a9d9852 f1ee247 89b9813 08c04fc 3208a74 a677593 3208a74 bfa85db 3208a74 33a9da7 bfa85db a677593 704eccb a9d9852 a3217ba f1ee247 f9c716f 48180e7 f9c716f 9df27df 2e08e3b 09a0506 3a00f55 a9363f4 5754a1c 313bb52 5754a1c 2e08e3b 09a0506 f4a1565 e17298c 09a0506 f9c716f a9d9852 e712a69 7ef36c6 fd35ae8 f1ee247 5754a1c 3a00f55 09a0506 50eea75 e712a69 7ef36c6 d0185ca 50eea75 704eccb 09a0506 704eccb 09a0506 704eccb 2e08e3b a9d9852 e137c08 b07d75e fe4f4a4 b07d75e fc0af52 5754a1c 53ee11a 1415561 f9c716f b403864 a677593 b07d75e a677593 599777e 8bb7641 9c3c2d7 89b9813 db037f3 89b9813 fc0af52 9df27df c575e18 8435838 fc0af52 1679b8f 8435838 f0231ae 1415561 ae12499 bc57790 e8716e2 23e2ef6 313bb52 98a4a00 1e27bfe 5845822 eac6a5b 5845822 313bb52 9c2b533 313bb52 98a4a00 c94bfc7 eac6a5b 0c3dd1c d249772 09eba18 8435838 e62ac65 8a73e0c 1679b8f e62ac65 4a67929 53ee11a a076741 f1ee247 054642a 313bb52 d249772 98a4a00 5845822 09eba18 a2d3df0 e62ac65 ec4fd5b eac6a5b e750aef 9df27df c144d28 0a01c1a 5845822 eac6a5b ec4fd5b 7ef36c6 c94bfc7 ec4fd5b d249772 e62ac65 a2d3df0 e62ac65 8435838 e62ac65 2548f77 8435838 e62ac65 2548f77 c575e18 8435838 1679b8f e62ac65 1679b8f 2548f77 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 |
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from typing import List, Tuple
import numpy as np
from PIL import Image, ImageDraw
import base64
import io
import asyncio
from utils import initialize_model, sample_frame
import torch
import os
import time
DEBUG = False
DEBUG_TEACHER_FORCING = False
app = FastAPI()
# Mount the static directory to serve HTML, JavaScript, and CSS files
app.mount("/static", StaticFiles(directory="static"), name="static")
# Add this at the top with other global variables
all_click_positions = [] # Store all historical click positions
def parse_action_string(action_str):
"""Convert formatted action string to x, y coordinates
Args:
action_str: String like 'N N N N N : N N N N N' or '+ 0 2 1 3 : + 0 3 8 3'
Returns:
tuple: (x, y) coordinates or None if action is padding
"""
action_type = action_str[0]
action_str = action_str[1:].strip()
if 'N' in action_str:
return (None, None, None)
# Split into x and y parts
action_str = action_str.replace(' ', '')
x_part, y_part = action_str.split(':')
# Parse x: remove sign, join digits, convert to int, apply sign
x = int(x_part)
# Parse y: remove sign, join digits, convert to int, apply sign
y = int(y_part)
return x, y, action_type
def create_position_and_click_map(pos,action_type, image_height=48, image_width=64, original_width=512, original_height=384):
"""Convert cursor position to a binary position map
Args:
x, y: Original cursor positions
image_size: Size of the output position map (square)
original_width: Original screen width (1024)
original_height: Original screen height (640)
Returns:
torch.Tensor: Binary position map of shape (1, image_size, image_size)
"""
x, y = pos
if x is None:
return torch.zeros((1, image_height, image_width)), torch.zeros((1, image_height, image_width)), None, None
# Scale the positions to new size
#x_scaled = int((x / original_width) * image_size)
#y_scaled = int((y / original_height) * image_size)
#screen_width, screen_height = 512, 384
#video_width, video_height = 512, 384
#x_scaled = x - (screen_width / 2 - video_width / 2)
#y_scaled = y - (screen_height / 2 - video_height / 2)
x_scaled = int(x / original_width * image_width)
y_scaled = int(y / original_height * image_height)
# Clamp values to ensure they're within bounds
x_scaled = max(0, min(x_scaled, image_width - 1))
y_scaled = max(0, min(y_scaled, image_height - 1))
# Create binary position map
pos_map = torch.zeros((1, image_height, image_width))
pos_map[0, y_scaled, x_scaled] = 1.0
leftclick_map = torch.zeros((1, image_height, image_width))
if action_type == 'L':
print ('left click', x_scaled, y_scaled)
print ('skipped')
if False:
leftclick_map[0, y_scaled, x_scaled] = 1.0
return pos_map, leftclick_map, x_scaled, y_scaled
# Serve the index.html file at the root URL
@app.get("/")
async def get():
return HTMLResponse(open("static/index.html").read())
def generate_random_image(width: int, height: int) -> np.ndarray:
return np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
def draw_trace(image: np.ndarray, previous_actions: List[Tuple[str, List[int]]], x_scaled=-1, y_scaled=-1) -> np.ndarray:
pil_image = Image.fromarray(image)
draw = ImageDraw.Draw(pil_image)
# Draw all historical click positions
for click_x, click_y in all_click_positions:
x_draw = click_x # Scale factor for display
y_draw = click_y
# Draw historical clicks as red circles
draw.ellipse([x_draw-4, y_draw-4, x_draw+4, y_draw+4], fill=(255, 0, 0))
# Draw current trace
prev_x, prev_y = None, None
for i, (action_type, position) in enumerate(previous_actions):
x, y = position
if x == 0 and y == 0:
continue
x_draw = x
y_draw = y
# Draw movement positions as blue dots
draw.ellipse([x_draw-2, y_draw-2, x_draw+2, y_draw+2], fill=(0, 0, 255))
# Draw connecting lines
if prev_x is not None:
draw.line([prev_x, prev_y, x_draw, y_draw], fill=(0, 255, 0), width=1)
prev_x, prev_y = x_draw, y_draw
# Draw current position
if x_scaled >= 0 and y_scaled >= 0:
x_current = x_scaled * 8
y_current = y_scaled * 8
draw.ellipse([x_current-3, y_current-3, x_current+3, y_current+3], fill=(0, 255, 0))
return np.array(pil_image)
# Initialize the model at the start of your application
#model = initialize_model("config_csllm.yaml", "yuntian-deng/computer-model")
model = initialize_model("pssearch_bsz64_acc1_lr8e5_512_leftclick_histpos_512_384.yaml", "yuntian-deng/computer-model")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
def load_initial_images(width, height):
initial_images = []
if DEBUG_TEACHER_FORCING:
# Load the previous 7 frames for image_81
for i in range(117-7, 117): # Load images 74-80
img = Image.open(f"record_10003/image_{i}.png")#.resize((width, height))
initial_images.append(np.array(img))
else:
#assert False
for i in range(7):
initial_images.append(np.zeros((height, width, 3), dtype=np.uint8))
return initial_images
def normalize_images(images, target_range=(-1, 1)):
images = np.stack(images).astype(np.float32)
if target_range == (-1, 1):
return images / 127.5 - 1
elif target_range == (0, 1):
return images / 255.0
else:
raise ValueError(f"Unsupported target range: {target_range}")
def denormalize_image(image, source_range=(-1, 1)):
if source_range == (-1, 1):
return ((image + 1) * 127.5).clip(0, 255).astype(np.uint8)
elif source_range == (0, 1):
return (image * 255).clip(0, 255).astype(np.uint8)
else:
raise ValueError(f"Unsupported source range: {source_range}")
def format_action(action_str, is_padding=False, is_leftclick=False):
if is_padding:
return "N N N N N N : N N N N N"
# Split the x~y coordinates
x, y = map(int, action_str.split('~'))
prefix = 'N'
if is_leftclick:
prefix = 'L'
# Convert numbers to padded strings and add spaces between digits
x_str = f"{abs(x):04d}"
y_str = f"{abs(y):04d}"
x_spaced = ' '.join(x_str)
y_spaced = ' '.join(y_str)
# Format with sign and proper spacing
return prefix + " " + f"{'+ ' if x >= 0 else '- '}{x_spaced} : {'+ ' if y >= 0 else '- '}{y_spaced}"
def predict_next_frame(previous_frames: List[np.ndarray], previous_actions: List[Tuple[str, List[int]]]) -> np.ndarray:
width, height = 512, 384
initial_images = load_initial_images(width, height)
# Prepare the image sequence for the model
image_sequence = previous_frames[-7:] # Take the last 7 frames
while len(image_sequence) < 7:
image_sequence.insert(0, initial_images[len(image_sequence)])
# Convert the image sequence to a tensor and concatenate in the channel dimension
image_sequence_tensor = torch.from_numpy(normalize_images(image_sequence, target_range=(-1, 1)))
image_sequence_tensor = image_sequence_tensor.to(device)
# Prepare the prompt based on the previous actions
action_descriptions = []
#initial_actions = ['901:604', '901:604', '901:604', '901:604', '901:604', '901:604', '901:604', '921:604']
initial_actions = ['0:0'] * 7
#initial_actions = ['N N N N N : N N N N N'] * 7
def unnorm_coords(x, y):
return int(x), int(y) #int(x - (1920 - 256) / 2), int(y - (1080 - 256) / 2)
# Process initial actions if there are not enough previous actions
while len(previous_actions) < 8:
x, y = map(int, initial_actions.pop(0).split(':'))
previous_actions.insert(0, ("N", unnorm_coords(x, y)))
prev_x = 0
prev_y = 0
#print ('here')
for action_type, pos in previous_actions: #[-8:]:
print ('here3', action_type, pos)
if action_type == 'move':
action_type = 'N'
if action_type == 'left_click':
action_type = 'L'
if action_type == "N":
x, y = pos
#norm_x = int(round(x / 256 * 1024)) #x + (1920 - 256) / 2
#norm_y = int(round(y / 256 * 640)) #y + (1080 - 256) / 2
#norm_x = x + (1920 - 512) / 2
#norm_y = y + (1080 - 512) / 2
norm_x = x
norm_y = y
if False and DEBUG_TEACHER_FORCING:
norm_x = x
norm_y = y
#action_descriptions.append(f"{(norm_x-prev_x):.0f}~{(norm_y-prev_y):.0f}")
#action_descriptions.append(format_action(f'{norm_x-prev_x:.0f}~{norm_y-prev_y:.0f}', x==0 and y==0))
action_descriptions.append(format_action(f'{norm_x:.0f}~{norm_y:.0f}', x==0 and y==0))
prev_x = norm_x
prev_y = norm_y
elif action_type == "L":
x, y = pos
#norm_x = int(round(x / 256 * 1024)) #x + (1920 - 256) / 2
#norm_y = int(round(y / 256 * 640)) #y + (1080 - 256) / 2
#norm_x = x + (1920 - 512) / 2
#norm_y = y + (1080 - 512) / 2
norm_x = x
norm_y = y
if False and DEBUG_TEACHER_FORCING:
norm_x = x #+ (1920 - 512) / 2
norm_y = y #+ (1080 - 512) / 2
#if DEBUG:
# norm_x = x
# norm_y = y
#action_descriptions.append(f"{(norm_x-prev_x):.0f}~{(norm_y-prev_y):.0f}")
#action_descriptions.append(format_action(f'{norm_x-prev_x:.0f}~{norm_y-prev_y:.0f}', x==0 and y==0))
action_descriptions.append(format_action(f'{norm_x:.0f}~{norm_y:.0f}', x==0 and y==0, True))
elif action_type == "right_click":
assert False
action_descriptions.append("right_click")
else:
assert False
prompt = " ".join(action_descriptions[-8:])
print(prompt)
#prompt = "N N N N N : N N N N N N N N N N : N N N N N N N N N N : N N N N N N N N N N : N N N N N N N N N N : N N N N N N N N N N : N N N N N N N N N N : N N N N N + 0 3 0 7 : + 0 3 7 5"
#x, y, action_type = parse_action_string(action_descriptions[-1])
#pos_map, leftclick_map, x_scaled, y_scaled = create_position_and_click_map((x, y), action_type)
leftclick_maps = []
pos_maps = []
for j in range(1, 9):
print ('fsfs', action_descriptions[-j])
x, y, action_type = parse_action_string(action_descriptions[-j])
pos_map_j, leftclick_map_j, x_scaled_j, y_scaled_j = create_position_and_click_map((x, y), action_type)
leftclick_maps.append(leftclick_map_j)
pos_maps.append(pos_map_j)
if j == 1:
x_scaled = x_scaled_j
y_scaled = y_scaled_j
if action_type == 'L':
all_click_positions.append((x, y))
#prompt = ''
#prompt = "1~1 0~0 0~0 0~0 0~0 0~0 0~0 0~0"
print(prompt)
#prompt = prompt.replace('L', 'N')
#print ('changing L to N')
# Generate the next frame
new_frame = sample_frame(model, prompt, image_sequence_tensor, pos_maps=pos_maps, leftclick_maps=leftclick_maps)
# Convert the generated frame to the correct format
new_frame = new_frame.transpose(1, 2, 0)
print (new_frame.max(), new_frame.min())
new_frame_denormalized = denormalize_image(new_frame, source_range=(-1, 1))
# Draw the trace of previous actions
new_frame_with_trace = draw_trace(new_frame_denormalized, previous_actions, x_scaled, y_scaled)
# Track click positions
#x, y, action_type = parse_action_string(action_descriptions[-1])
return new_frame_with_trace, new_frame_denormalized
# WebSocket endpoint for continuous user interaction
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
global all_click_positions # Add this line
all_click_positions = [] # Reset at the start of each connection
client_id = id(websocket) # Use a unique identifier for each connection
print(f"New WebSocket connection: {client_id}")
await websocket.accept()
previous_frames = []
previous_actions = []
positions = ['815~335', '787~342', '787~342', '749~345', '703~346', '703~346', '654~347', '654~347', '604~349', '555~353', '555~353', '509~357', '509~357', '468~362', '431~368', '431~368']
#positions = ['815~335', '787~342', '749~345', '703~346', '703~346', '654~347', '654~347', '604~349', '555~353', '555~353', '509~357', '509~357', '468~362', '431~368', '431~368']
positions = ['307~375']
positions = ['815~335']
#positions = ['787~342']
positions = ['300~800']
if DEBUG_TEACHER_FORCING:
#print ('here2')
# Use the predefined actions for image_81
debug_actions = [
'N + 0 8 5 3 : + 0 4 5 0', 'N + 0 8 7 1 : + 0 4 6 3',
'N + 0 8 9 0 : + 0 4 7 5', 'N + 0 9 0 8 : + 0 4 8 8',
'N + 0 9 2 7 : + 0 5 0 1', 'N + 0 9 2 7 : + 0 5 0 1',
'N + 0 9 2 7 : + 0 5 0 1', 'N + 0 9 2 7 : + 0 5 0 1',
'N + 0 9 2 7 : + 0 5 0 1', 'N + 0 9 2 7 : + 0 5 0 1',
'L + 0 9 2 7 : + 0 5 0 1', 'N + 0 9 2 7 : + 0 5 0 1',
'L + 0 9 2 7 : + 0 5 0 1', 'N + 0 9 2 7 : + 0 5 0 1',
'N + 0 9 2 7 : + 0 5 0 1', #'N + 0 9 2 7 : + 0 5 0 1'
]
debug_actions = [
'N + 1 1 6 5 : + 0 4 4 3', 'N + 1 1 7 0 : + 0 4 1 8',
'N + 1 1 7 5 : + 0 3 9 4', 'N + 1 1 8 1 : + 0 3 7 0',
'N + 1 1 8 4 : + 0 3 5 8', 'N + 1 1 8 9 : + 0 3 3 3',
'N + 1 1 9 4 : + 0 3 0 9', 'N + 1 1 9 7 : + 0 2 9 7',
'N + 1 1 9 7 : + 0 2 9 7', 'N + 1 1 9 7 : + 0 2 9 7',
'N + 1 1 9 7 : + 0 2 9 7', 'N + 1 1 9 7 : + 0 2 9 7',
'L + 1 1 9 7 : + 0 2 9 7', 'N + 1 1 9 7 : + 0 2 9 7',
'N + 1 1 9 7 : + 0 2 9 7'
]
debug_actions = [
'N + 1 1 6 5 : + 0 4 4 3', 'N + 1 1 7 0 : + 0 4 1 8',
'N + 1 1 7 5 : + 0 3 9 4', 'N + 1 1 8 1 : + 0 3 7 0',
'N + 1 1 8 4 : + 0 3 5 8', 'N + 1 1 8 9 : + 0 3 3 3',
'N + 1 1 9 4 : + 0 3 0 9', 'N + 1 1 9 7 : + 0 2 9 7',
'N + 1 1 9 7 : + 0 2 9 7', 'N + 1 1 9 7 : + 0 2 9 7',
'N + 1 1 9 7 : + 0 2 9 7', 'N + 1 1 9 7 : + 0 2 9 7',
'N + 1 1 9 7 : + 0 2 9 7', 'N + 1 1 9 7 : + 0 2 9 7',
'N + 1 1 9 7 : + 0 2 9 7'
]
debug_actions = ['N + 0 0 4 0 : + 0 2 0 4', 'N + 0 1 3 8 : + 0 1 9 0',
'N + 0 2 7 4 : + 0 3 8 3', 'N + 0 5 0 1 : + 0 1 7 3',
'L + 0 4 7 3 : + 0 0 8 7', 'N + 0 1 0 9 : + 0 3 4 4',
'N + 0 0 5 2 : + 0 1 9 4', 'N + 0 3 6 5 : + 0 2 3 2',
'N + 0 3 8 9 : + 0 2 4 5', 'N + 0 0 2 0 : + 0 0 5 9',
'N + 0 4 7 3 : + 0 1 5 7', 'L + 0 1 9 1 : + 0 0 8 7',
'L + 0 1 9 1 : + 0 0 8 7', 'N + 0 3 4 3 : + 0 2 6 3', ]
#'N + 0 2 0 5 : + 0 1 3 3']
previous_actions = []
for action in debug_actions[-8:]:
action = action.replace('1 1', '0 4')
x, y, action_type = parse_action_string(action)
previous_actions.append((action_type, (x, y)))
positions = [
'N + 0 9 2 7 : + 0 5 0 1', 'N + 0 9 1 8 : + 0 4 9 2',
'N + 0 9 0 8 : + 0 4 8 3', 'N + 0 8 9 8 : + 0 4 7 4',
'N + 0 8 8 9 : + 0 4 6 5', 'N + 0 8 8 0 : + 0 4 5 6',
'N + 0 8 7 0 : + 0 4 4 7', 'N + 0 8 6 0 : + 0 4 3 8',
'N + 0 8 5 1 : + 0 4 2 9', 'N + 0 8 4 2 : + 0 4 2 0',
'N + 0 8 3 2 : + 0 4 1 1', 'N + 0 8 3 2 : + 0 4 1 1'
]
positions = [
#'L + 1 1 9 7 : + 0 2 9 7', 'N + 1 1 9 7 : + 0 2 9 7',
'N + 1 1 9 7 : + 0 2 9 7', 'N + 1 1 9 7 : + 0 2 9 7',
'N + 1 1 7 9 : + 0 3 0 3', 'N + 1 1 4 2 : + 0 3 1 4',
'N + 1 1 0 6 : + 0 3 2 6', 'N + 1 0 6 9 : + 0 3 3 7',
'N + 1 0 5 1 : + 0 3 4 3', 'N + 1 0 1 4 : + 0 3 5 4',
'N + 0 9 7 8 : + 0 3 6 5', 'N + 0 9 4 2 : + 0 3 7 7',
'N + 0 9 0 5 : + 0 3 8 8', 'N + 0 8 6 8 : + 0 4 0 0',
'N + 0 8 3 2 : + 0 4 1 1'
]
positions = ['L + 0 1 9 1 : + 0 0 8 7',
'L + 0 1 9 1 : + 0 0 8 7', 'N + 0 3 4 3 : + 0 2 6 3',
'N + 0 2 0 5 : + 0 1 3 3', 'N + 0 0 7 6 : + 0 3 4 5',
'N + 0 3 1 8 : + 0 3 3 3', 'N + 0 2 5 4 : + 0 2 9 0',
'N + 0 1 0 6 : + 0 1 6 4', 'N + 0 0 7 4 : + 0 2 8 4',
'N + 0 0 2 4 : + 0 0 4 1', 'N + 0 1 5 0 : + 0 3 8 3',
'N + 0 4 0 5 : + 0 1 6 8', 'N + 0 0 5 4 : + 0 3 2 4',
'N + 0 2 9 0 : + 0 1 4 1', 'N + 0 4 0 2 : + 0 0 0 9',
'N + 0 3 0 7 : + 0 3 3 2', 'N + 0 2 2 0 : + 0 3 7 1',
'N + 0 0 8 2 : + 0 1 5 1']
positions = positions[3:]
#positions = positions[:4]
position = positions[0]
positions = positions[1:]
x, y, action_type = parse_action_string(position)
mouse_position = (x, y)
previous_actions.append((action_type, mouse_position))
if not DEBUG_TEACHER_FORCING:
previous_actions = []
try:
while True:
try:
# Receive user input with a timeout
#data = await asyncio.wait_for(websocket.receive_json(), timeout=90000.0)
data = await websocket.receive_json()
if data.get("type") == "heartbeat":
await websocket.send_json({"type": "heartbeat_response"})
continue
action_type = data.get("action_type")
mouse_position = data.get("mouse_position")
# Store the actions
if DEBUG:
position = positions[0]
#positions = positions[1:]
#mouse_position = position.split('~')
#mouse_position = [int(item) for item in mouse_position]
#mouse_position = '+ 0 8 1 5 : + 0 3 3 5'
if False and DEBUG_TEACHER_FORCING:
position = positions[0]
positions = positions[1:]
x, y, action_type = parse_action_string(position)
mouse_position = (x, y)
if False:
previous_actions.append((action_type, mouse_position))
#previous_actions = [(action_type, mouse_position)]
if not DEBUG_TEACHER_FORCING:
previous_actions.append((action_type, mouse_position))
# Log the start time
start_time = time.time()
# Predict the next frame based on the previous frames and actions
if DEBUG_TEACHER_FORCING:
print ('predicting', f"record_10003/image_{117+len(previous_frames)}.png")
next_frame, next_frame_append = predict_next_frame(previous_frames, previous_actions)
# Load and append the corresponding ground truth image instead of model output
print ('here4', len(previous_frames))
if True and DEBUG_TEACHER_FORCING:
img = Image.open(f"record_10003/image_{117+len(previous_frames)}.png")
previous_frames.append(img)
elif True:
assert False
previous_frames.append(next_frame_append)
previous_frames = []
# Convert the numpy array to a base64 encoded image
img = Image.fromarray(next_frame)
buffered = io.BytesIO()
img.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode()
# Log the processing time
processing_time = time.time() - start_time
print(f"Frame processing time: {processing_time:.2f} seconds")
# Send the generated frame back to the client
await websocket.send_json({"image": img_str})
except asyncio.TimeoutError:
print("WebSocket connection timed out")
#break # Exit the loop on timeout
except WebSocketDisconnect:
print("WebSocket disconnected")
#break # Exit the loop on disconnect
except Exception as e:
print(f"Error in WebSocket connection {client_id}: {e}")
finally:
print(f"WebSocket connection closed: {client_id}")
#await websocket.close() # Ensure the WebSocket is closed
|