alphard/efficiency-nodes-comfyui
1
0
Fork 0
mirror of https://github.com/jags111/efficiency-nodes-comfyui.git synced 2026-03-25 15:15:45 -03:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #93 from LucianoCirino/Improved-server_logic-WebSocket-Handling

Browse Source 
Improved server_logic WebSocket Handling
This commit is contained in:
TSC
2023-08-10 19:51:18 -05:00
committed by GitHub
parent 60e0b1bcc0 235e6aea01
commit daf5645ee6
1 changed files with 17 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
tsc_utils.py
View File

@@ -524,7 +524,7 @@ import base64
from io import BytesIO from io import BytesIO
from torchvision import transforms from torchvision import transforms
latest_image = None latest_image = list()
connected_client = None connected_client = None
websocket_status = True websocket_status = True
@@ -538,17 +538,26 @@ def handle_websocket_failure():
f"preview is enabled and vae decoding is set to 'true`.") f"preview is enabled and vae decoding is set to 'true`.")
async def server_logic(websocket, path): async def server_logic(websocket, path):
global latest_image, connected_client global latest_image, connected_client, websocket_status
# If websocket_status is False, set latest_image to an empty list
if not websocket_status:
latest_image = list()
# Assign the connected client # Assign the connected client
connected_client = websocket connected_client = websocket
async for message in websocket: try:
# If not a command, treat it as image data async for message in websocket:
if not message.startswith('{'): # If not a command, treat it as image data
image_data = base64.b64decode(message.split(",")[1]) if not message.startswith('{'):
image = Image.open(BytesIO(image_data)) image_data = base64.b64decode(message.split(",")[1])
latest_image = pil2tensor(image) image = Image.open(BytesIO(image_data))
latest_image = pil2tensor(image)
except (websockets.exceptions.ConnectionClosedError, asyncio.exceptions.CancelledError):
handle_websocket_failure()
except Exception:
handle_websocket_failure()
def run_server(): def run_server():
loop = asyncio.new_event_loop() loop = asyncio.new_event_loop()