From 286677e5d4aec442dc3d9cabcb42f61c8351ccc8 Mon Sep 17 00:00:00 2001
From: e-maks <my@mmatantsev.ru>
Date: Sat, 7 Mar 2026 21:59:43 +0300
Subject: [PATCH] handle errors

---
 server.py | 53 +++++++++++++++++++++++++++++++++++------------------
 1 file changed, 35 insertions(+), 18 deletions(-)

diff --git a/server.py b/server.py
index 0303a5d..12cb4b3 100644
--- a/server.py
+++ b/server.py
@@ -159,24 +159,41 @@ def generate_video(prompt):
     ws_url = f"ws://{COMFY_HOST}:{COMFY_PORT}/ws?clientId={ws_client_id}"
 
     ws = ws_client.WebSocket()
-    ws.connect(ws_url)
-    logger.info("🔌 WebSocket подключён к ComfyUI")
-
-    prompt_id = queue_prompt(prompt)["prompt_id"]
-    logger.info(f"📤 Prompt отправлен: {prompt_id}")
-
-    # Ждём завершения
-    while True:
-        out = ws.recv()
-        if isinstance(out, str):
-            message = json.loads(out)
-            if message["type"] == "executing":
-                data = message["data"]
-                if data["node"] is None and data["prompt_id"] == prompt_id:
-                    break
-
-    ws.close()
-    logger.info("✅ Генерация завершена")
+    try:
+        ws.connect(ws_url)
+        logger.info("🔌 WebSocket подключён к ComfyUI")
+
+        prompt_id = queue_prompt(prompt)["prompt_id"]
+        logger.info(f"📤 Prompt отправлен: {prompt_id}")
+
+        # Ждём завершения
+        while True:
+            try:
+                out = ws.recv()
+            except ws_client.WebSocketConnectionClosedException:
+                # ComfyUI разорвал соединение — проверяем что случилось
+                logger.error(f"⚠️ ComfyUI закрыл WebSocket для prompt {prompt_id}")
+                raise RuntimeError(
+                    "ComfyUI connection lost during generation. "
+                    "Possible causes: Out of VRAM, workflow error, or ComfyUI crash. "
+                    "Check ComfyUI logs: journalctl -u comfyui -n 100"
+                )
+            
+            if isinstance(out, str):
+                message = json.loads(out)
+                if message["type"] == "executing":
+                    data = message["data"]
+                    if data["node"] is None and data["prompt_id"] == prompt_id:
+                        break
+
+        ws.close()
+        logger.info("✅ Генерация завершена")
+    except Exception as e:
+        try:
+            ws.close()
+        except:
+            pass
+        raise
 
     # Извлекаем видео
     history = get_history(prompt_id)[prompt_id]