quant-speed-AI/V2_micropython
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

1

Browse Source
This commit is contained in:
jeremygan2021
2026-03-03 23:07:17 +08:00
parent 05f02a1454
commit 0aa8f5f473
7 changed files with 204 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

142
main.py
View File

@@ -23,6 +23,7 @@ SERVER_URL = f"ws://{SERVER_IP}:{SERVER_PORT}/ws/audio"
IMAGE_STATE_IDLE = 0
IMAGE_STATE_RECEIVING = 1
UI_SCREEN_HOME = 0
UI_SCREEN_RECORDING = 1
UI_SCREEN_CONFIRM = 2
UI_SCREEN_RESULT = 3
@@ -39,7 +40,7 @@ _btn_release_time = 0
_btn_press_time = 0
def connect_wifi(max_retries=5):
def connect_wifi(display=None, max_retries=5):
"""连接WiFi网络"""
wlan = network.WLAN(network.STA_IF)
@@ -50,33 +51,53 @@ def connect_wifi(max_retries=5):
time.sleep(3)
except Exception as e:
print(f"WiFi init error: {e}")
if display and display.tft:
display.render_wifi_status(False)
return False
if display and display.tft:
display.render_wifi_connecting()
for attempt in range(max_retries):
try:
if wlan.isconnected():
print('WiFi connected')
if display and display.tft:
display.render_wifi_status(True)
time.sleep(1.5)
return True
print(f'Connecting to WiFi {WIFI_SSID}...')
wlan.connect(WIFI_SSID, WIFI_PASS)
start_time = time.ticks_ms()
spinner_angle = 0
while not wlan.isconnected():
if time.ticks_diff(time.ticks_ms(), start_time) > 30000:
print("WiFi timeout!")
break
time.sleep(0.5)
time.sleep(0.1)
print(".", end="")
# 简单的加载动画
if display and display.tft:
if time.ticks_ms() % 200 < 50: # 节流刷新
draw_loading_spinner(display, 120, 150, spinner_angle, st7789.CYAN)
spinner_angle = (spinner_angle + 45) % 360
if wlan.isconnected():
print('\nWiFi connected!')
if display and display.tft:
display.render_wifi_status(True)
time.sleep(1.5)
return True
if attempt < max_retries - 1:
print(f"\nRetry {attempt + 1}/{max_retries}...")
wlan.disconnect()
time.sleep(3)
if display and display.tft:
display.text(f"重试 {attempt + 1}/{max_retries}...", 80, 180, st7789.YELLOW, wait=False)
except Exception as e:
print(f"WiFi error: {e}")
@@ -84,6 +105,9 @@ def connect_wifi(max_retries=5):
time.sleep(5)
print("WiFi connection failed!")
if display and display.tft:
display.render_wifi_status(False)
time.sleep(3)
return False
@@ -216,9 +240,10 @@ def render_result_screen(display, status="", prompt="", image_received=False):
elif status == "COMPLETE" or image_received:
# Don't clear screen, image is already there
# display.text("生成完成!", 80, 50, st7789.GREEN)
# draw_check_icon(display, 110, 80)
pass
# Draw a small indicator to show it's done, but don't cover the image
# Maybe a small green dot in the corner?
display.tft.fill_rect(230, 230, 10, 10, st7789.GREEN)
elif status == "ERROR":
display.tft.fill(st7789.BLACK)
@@ -226,14 +251,14 @@ def render_result_screen(display, status="", prompt="", image_received=False):
display.text("AI 生成中", 80, 8, st7789.BLACK)
display.text("生成失败", 80, 50, st7789.RED)
if prompt and not image_received:
if prompt and not image_received and not image_generation_done:
display.tft.fill_rect(10, 140, 220, 50, 0x2124) # Dark Grey
display.text("提示词:", 15, 145, st7789.CYAN)
display.text(prompt[:25] + "..." if len(prompt) > 25 else prompt, 15, 165, st7789.WHITE)
# Only show back button if not showing full image, or maybe show it transparently?
# For now, let's not cover the image with the button hint
if not image_received:
if not image_received and not image_generation_done:
display.tft.fill_rect(60, 210, 120, 25, st7789.BLUE)
display.text("长按返回", 90, 215, st7789.WHITE)
@@ -421,8 +446,10 @@ def main():
if display.tft:
display.init_ui()
display.render_home_screen()
time.sleep(2)
ui_screen = UI_SCREEN_RECORDING
ui_screen = UI_SCREEN_HOME
is_recording = False
ws = None
image_state = IMAGE_STATE_IDLE
@@ -448,23 +475,40 @@ def main():
while retry_count < max_retries:
try:
print(f"Connecting to {SERVER_URL} (attempt {retry_count + 1})")
if display and display.tft:
display.tft.fill_rect(0, 220, 240, 20, st7789.BLACK)
display.text(f"连接服务器...({retry_count+1})", 60, 220, st7789.CYAN, wait=False)
ws = WebSocketClient(SERVER_URL)
print("WebSocket connected!")
if display:
display.set_ws(ws)
# 预热字体，请求常用字
# 可以在这里发一个 GET_HIGH_FREQ 请求，或者简单的不做处理，因为 render_home_screen 已经触发了部分
return True
except Exception as e:
print(f"WS connection failed: {e}")
retry_count += 1
time.sleep(1)
if display and display.tft:
display.text("服务器连接失败", 60, 220, st7789.RED, wait=False)
time.sleep(2)
return False
if connect_wifi():
if connect_wifi(display):
connect_ws()
# WiFi 和 WS 都连接成功后，进入录音界面
ui_screen = UI_SCREEN_RECORDING
if display.tft:
render_recording_screen(display, "", 0)
else:
print("Running in offline mode")
# 即使离线也进入录音界面（虽然不能用）
ui_screen = UI_SCREEN_RECORDING
if display.tft:
render_recording_screen(display, "离线模式", 0)
read_buf = bytearray(4096)
last_audio_level = 0
@@ -542,9 +586,12 @@ def main():
elif ui_screen == UI_SCREEN_CONFIRM:
print(">>> Confirm and generate")
# 发送生成图片指令
if ws and ws.is_connected():
try:
ws.send("STOP_RECORDING")
# 明确发送生成指令
ws.send(f"GENERATE_IMAGE:{current_asr_text}")
except:
ws = None
@@ -575,7 +622,9 @@ def main():
is_recording = False
if ui_screen == UI_SCREEN_RECORDING or is_recording == False:
# If in recording screen or (not recording AND not result screen), then regenerate/re-record
# This ensures result screen is handled by its own block below
if ui_screen == UI_SCREEN_RECORDING:
if current_asr_text:
print(">>> Generate image with ASR text")
ui_screen = UI_SCREEN_RESULT
@@ -634,40 +683,41 @@ def main():
try:
ws.send(read_buf[:num_read], opcode=2)
poller = uselect.poll()
poller.register(ws.sock, uselect.POLLIN)
events = poller.poll(0)
if events:
msg = ws.recv()
image_state, event_data = process_message(msg, display, image_state, image_data_list)
if event_data:
if event_data[0] == "asr":
current_asr_text = event_data[1]
if display.tft:
render_recording_screen(display, current_asr_text, last_audio_level)
elif event_data[0] == "font_update":
if ui_screen == UI_SCREEN_RECORDING and display.tft:
render_recording_screen(display, current_asr_text, last_audio_level)
elif event_data[0] == "status":
current_status = event_data[1]
status_text = event_data[2] if len(event_data) > 2 else ""
if display.tft:
render_result_screen(display, current_status, current_prompt, image_generation_done)
elif event_data[0] == "prompt":
current_prompt = event_data[1]
elif event_data[0] == "image_done":
image_generation_done = True
if display.tft:
render_result_screen(display, "COMPLETE", current_prompt, True)
elif event_data[0] == "error":
if display.tft:
render_result_screen(display, "ERROR", current_prompt, False)
# 移除录音时的消息接收，确保录音流畅
# poller = uselect.poll()
# poller.register(ws.sock, uselect.POLLIN)
# events = poller.poll(0)
# if events:
# msg = ws.recv()
# image_state, event_data = process_message(msg, display, image_state, image_data_list)
#
# if event_data:
# if event_data[0] == "asr":
# current_asr_text = event_data[1]
# if display.tft:
# render_recording_screen(display, current_asr_text, last_audio_level)
#
# elif event_data[0] == "font_update":
# if ui_screen == UI_SCREEN_RECORDING and display.tft:
# render_recording_screen(display, current_asr_text, last_audio_level)
#
# elif event_data[0] == "status":
# current_status = event_data[1]
# status_text = event_data[2] if len(event_data) > 2 else ""
# if display.tft:
# render_result_screen(display, current_status, current_prompt, image_generation_done)
#
# elif event_data[0] == "prompt":
# current_prompt = event_data[1]
#
# elif event_data[0] == "image_done":
# image_generation_done = True
# if display.tft:
# render_result_screen(display, "COMPLETE", current_prompt, True)
#
# elif event_data[0] == "error":
# if display.tft:
# render_result_screen(display, "ERROR", current_prompt, False)
except:
ws = None