refactor: fix websocket errors

frontend/src/lib/components/ImagePlayer.svelte

@@ -7,8 +7,8 @@
   import Expand from '$lib/icons/expand.svelte';
   import { snapImage, expandWindow } from '$lib/utils';
 
-  $: isLCMRunning = $lcmLiveStatus !== LCMLiveStatus.DISCONNECTED;
-  $: console.log('isLCMRunning', isLCMRunning);
+  $: isLCMRunning = $lcmLiveStatus !== LCMLiveStatus.DISCONNECTED && 
+                 $lcmLiveStatus !== LCMLiveStatus.ERROR;
   let imageEl: HTMLImageElement;
   let expandedWindow: Window;
   let isExpanded = false;
@@ -40,12 +40,26 @@
   class="relative mx-auto aspect-square max-w-lg self-center overflow-hidden rounded-lg border border-slate-300"
 >
   <!-- svelte-ignore a11y-missing-attribute -->
-  {#if isLCMRunning}
+  {#if $lcmLiveStatus === LCMLiveStatus.CONNECTING}
+    <!-- Show connecting spinner -->
+    <div class="flex items-center justify-center h-full w-full">
+      <div class="animate-spin rounded-full h-16 w-16 border-b-2 border-white"></div>
+      <p class="text-white ml-2">Connecting...</p>
+    </div>
+  {:else if isLCMRunning}
     {#if !isExpanded}
+      <!-- Handle image error by adding onerror event -->
       <img
         bind:this={imageEl}
         class="aspect-square w-full rounded-lg"
         src={'/api/stream/' + $streamId}
+        on:error={(e) => {
+          console.error('Image stream error:', e);
+          // If stream fails to load, set status to error
+          if ($lcmLiveStatus !== LCMLiveStatus.ERROR) {
+            lcmLiveStatus.set(LCMLiveStatus.ERROR);
+          }
+        }}
       />
     {/if}
     <div class="absolute bottom-1 right-1">
@@ -65,6 +79,13 @@
         <Floppy classList={''} />
       </Button>
     </div>
+  {:else if $lcmLiveStatus === LCMLiveStatus.ERROR}
+    <!-- Show error state with red border -->
+    <div class="flex items-center justify-center h-full w-full border-2 border-red-500 rounded-lg bg-gray-900">
+      <p class="text-center text-white p-4">
+        Connection error
+      </p>
+    </div>
   {:else}
     <img
       class="aspect-square w-full rounded-lg"

frontend/src/lib/lcmLive.ts

@@ -1,11 +1,13 @@
-import { writable } from 'svelte/store';
+import { get, writable } from 'svelte/store';
 
 export enum LCMLiveStatus {
   CONNECTED = 'connected',
   DISCONNECTED = 'disconnected',
+  CONNECTING = 'connecting',
   WAIT = 'wait',
   SEND_FRAME = 'send_frame',
-  TIMEOUT = 'timeout'
+  TIMEOUT = 'timeout',
+  ERROR = 'error'
 }
 
 const initStatus: LCMLiveStatus = LCMLiveStatus.DISCONNECTED;
@@ -13,85 +15,194 @@ const initStatus: LCMLiveStatus = LCMLiveStatus.DISCONNECTED;
 export const lcmLiveStatus = writable<LCMLiveStatus>(initStatus);
 export const streamId = writable<string | null>(null);
 
+// WebSocket connection
 let websocket: WebSocket | null = null;
+// Flag to track intentional connection closure
+let intentionalClosure = false;
+
+// Register browser unload event listener to properly close WebSockets
+if (typeof window !== 'undefined') {
+  window.addEventListener('beforeunload', () => {
+    // Mark any closure during page unload as intentional
+    intentionalClosure = true;
+    // Close the WebSocket properly if it exists
+    if (websocket && websocket.readyState === WebSocket.OPEN) {
+      websocket.close(1000, 'Page unload');
+    }
+  });
+}
 export const lcmLiveActions = {
   async start(getSreamdata: () => any[]) {
     return new Promise((resolve, reject) => {
       try {
+        // Set connecting status immediately
+        lcmLiveStatus.set(LCMLiveStatus.CONNECTING);
+
         const userId = crypto.randomUUID();
         const websocketURL = `${
           window.location.protocol === 'https:' ? 'wss' : 'ws'
         }:${window.location.host}/api/ws/${userId}`;
 
+        // Close any existing connection first
+        if (websocket && websocket.readyState !== WebSocket.CLOSED) {
+          websocket.close();
+        }
+
         websocket = new WebSocket(websocketURL);
+
+        // Set a connection timeout
+        const connectionTimeout = setTimeout(() => {
+          if (websocket && websocket.readyState !== WebSocket.OPEN) {
+            console.error('WebSocket connection timeout');
+            lcmLiveStatus.set(LCMLiveStatus.ERROR);
+            streamId.set(null);
+            reject(new Error('Connection timeout. Please try again.'));
+            websocket.close();
+          }
+        }, 10000); // 10 second timeout
+
         websocket.onopen = () => {
+          clearTimeout(connectionTimeout);
           console.log('Connected to websocket');
         };
-        websocket.onclose = () => {
-          lcmLiveStatus.set(LCMLiveStatus.DISCONNECTED);
-          console.log('Disconnected from websocket');
+
+        websocket.onclose = (event) => {
+          clearTimeout(connectionTimeout);
+          console.log(`Disconnected from websocket: ${event.code} ${event.reason}`);
+
+          // Only change status if we're not in ERROR state (which would mean we already handled the error)
+          if (get(lcmLiveStatus) !== LCMLiveStatus.ERROR) {
+            lcmLiveStatus.set(LCMLiveStatus.DISCONNECTED);
+          }
+
+          // If connection was never established (close without open)
+          if (event.code === 1006 && streamId.get() === null) {
+            reject(new Error('Cannot connect to server. Please try again later.'));
+          }
         };
+
         websocket.onerror = (err) => {
-          console.error(err);
+          clearTimeout(connectionTimeout);
+          console.error('WebSocket error:', err);
+          lcmLiveStatus.set(LCMLiveStatus.ERROR);
+          streamId.set(null);
+          reject(new Error('Connection error. Please try again.'));
         };
+
         websocket.onmessage = (event) => {
-          const data = JSON.parse(event.data);
-          switch (data.status) {
-            case 'connected':
-              lcmLiveStatus.set(LCMLiveStatus.CONNECTED);
-              streamId.set(userId);
-              resolve({ status: 'connected', userId });
-              break;
-            case 'send_frame':
-              lcmLiveStatus.set(LCMLiveStatus.SEND_FRAME);
-              const streamData = getSreamdata();
-              websocket?.send(JSON.stringify({ status: 'next_frame' }));
-              for (const d of streamData) {
-                this.send(d);
-              }
-              break;
-            case 'wait':
-              lcmLiveStatus.set(LCMLiveStatus.WAIT);
-              break;
-            case 'timeout':
-              console.log('timeout');
-              lcmLiveStatus.set(LCMLiveStatus.TIMEOUT);
-              streamId.set(null);
-              reject(new Error('timeout'));
-              break;
-            case 'error':
-              console.log(data.message);
-              lcmLiveStatus.set(LCMLiveStatus.DISCONNECTED);
-              streamId.set(null);
-              reject(new Error(data.message));
-              break;
+          try {
+            const data = JSON.parse(event.data);
+            switch (data.status) {
+              case 'connected':
+                lcmLiveStatus.set(LCMLiveStatus.CONNECTED);
+                streamId.set(userId);
+                resolve({ status: 'connected', userId });
+                break;
+              case 'send_frame':
+                lcmLiveStatus.set(LCMLiveStatus.SEND_FRAME);
+                try {
+                  const streamData = getSreamdata();
+                  // Send as an object, not a string, to use the proper handling in the send method
+                  this.send({ status: 'next_frame' });
+                  for (const d of streamData) {
+                    this.send(d);
+                  }
+                } catch (error) {
+                  console.error('Error sending frame data:', error);
+                }
+                break;
+              case 'wait':
+                lcmLiveStatus.set(LCMLiveStatus.WAIT);
+                break;
+              case 'timeout':
+                console.log('Session timeout');
+                lcmLiveStatus.set(LCMLiveStatus.TIMEOUT);
+                streamId.set(null);
+                reject(new Error('Session timeout. Please restart.'));
+                break;
+              case 'error':
+                console.error('Server error:', data.message);
+                lcmLiveStatus.set(LCMLiveStatus.ERROR);
+                streamId.set(null);
+                reject(new Error(data.message || 'Server error occurred'));
+                break;
+              default:
+                console.log('Unknown message status:', data.status);
+            }
+          } catch (error) {
+            console.error('Error handling websocket message:', error);
           }
         };
       } catch (err) {
-        console.error(err);
-        lcmLiveStatus.set(LCMLiveStatus.DISCONNECTED);
+        console.error('Error initializing websocket:', err);
+        lcmLiveStatus.set(LCMLiveStatus.ERROR);
         streamId.set(null);
         reject(err);
       }
     });
   },
   send(data: Blob | { [key: string]: any }) {
-    if (websocket && websocket.readyState === WebSocket.OPEN) {
-      if (data instanceof Blob) {
-        websocket.send(data);
+    try {
+      if (websocket && websocket.readyState === WebSocket.OPEN) {
+        if (data instanceof Blob) {
+          websocket.send(data);
+        } else {
+          websocket.send(JSON.stringify(data));
+        }
       } else {
-        websocket.send(JSON.stringify(data));
+        const readyStateText = websocket
+          ? ['CONNECTING', 'OPEN', 'CLOSING', 'CLOSED'][websocket.readyState]
+          : 'null';
+        console.warn(`WebSocket not ready for sending: ${readyStateText}`);
+
+        // If WebSocket is closed unexpectedly, set status to disconnected
+        if (!websocket || websocket.readyState === WebSocket.CLOSED) {
+          lcmLiveStatus.set(LCMLiveStatus.DISCONNECTED);
+          streamId.set(null);
+        }
       }
-    } else {
-      console.log('WebSocket not connected');
+    } catch (error) {
+      console.error('Error sending data through WebSocket:', error);
+      // Handle WebSocket error by forcing disconnection
+      this.stop();
     }
   },
+
+  async reconnect(getSreamdata: () => any[]) {
+    try {
+      await this.stop();
+      // Small delay to ensure clean disconnection before reconnecting
+      await new Promise((resolve) => setTimeout(resolve, 500));
+      return await this.start(getSreamdata);
+    } catch (error) {
+      console.error('Reconnection failed:', error);
+      throw error;
+    }
+  },
+
   async stop() {
     lcmLiveStatus.set(LCMLiveStatus.DISCONNECTED);
-    if (websocket) {
-      websocket.close();
+    try {
+      if (websocket) {
+        // Only attempt to close if not already closed
+        if (websocket.readyState !== WebSocket.CLOSED) {
+          // Set up onclose handler to clean up only
+          websocket.onclose = () => {
+            console.log('WebSocket closed cleanly during stop()');
+          };
+
+          // Set up onerror to be silent during intentional closure
+          websocket.onerror = () => {};
+
+          websocket.close(1000, 'Client initiated disconnect');
+        }
+      }
+    } catch (error) {
+      console.error('Error during WebSocket closure:', error);
+    } finally {
+      // Always clean up references
+      websocket = null;
+      streamId.set(null);
     }
-    websocket = null;
-    streamId.set(null);
   }
 };

frontend/src/routes/+page.svelte

@@ -57,35 +57,86 @@
     }
   }
 
-  $: isLCMRunning = $lcmLiveStatus !== LCMLiveStatus.DISCONNECTED;
-  $: if ($lcmLiveStatus === LCMLiveStatus.TIMEOUT) {
-    warningMessage = 'Session timed out. Please try again.';
+  $: isLCMRunning = $lcmLiveStatus !== LCMLiveStatus.DISCONNECTED && 
+                 $lcmLiveStatus !== LCMLiveStatus.ERROR;
+  $: isConnecting = $lcmLiveStatus === LCMLiveStatus.CONNECTING;
+                 
+  $: {
+    // Set warning messages based on lcmLiveStatus
+    if ($lcmLiveStatus === LCMLiveStatus.TIMEOUT) {
+      warningMessage = 'Session timed out. Please try again.';
+    } else if ($lcmLiveStatus === LCMLiveStatus.ERROR) {
+      warningMessage = 'Connection error occurred. Please try again.';
+    }
   }
   let disabled = false;
   async function toggleLcmLive() {
     try {
       if (!isLCMRunning) {
-        if (isImageMode) {
-          await mediaStreamActions.enumerateDevices();
-          await mediaStreamActions.start();
+        if (isConnecting) {
+          return; // Don't allow multiple connection attempts
         }
+        
+        // Clear any previous warning messages
+        warningMessage = '';
         disabled = true;
-        await lcmLiveActions.start(getSreamdata);
-        disabled = false;
-        toggleQueueChecker(false);
+        
+        try {
+          if (isImageMode) {
+            await mediaStreamActions.enumerateDevices();
+            await mediaStreamActions.start();
+          }
+          
+          await lcmLiveActions.start(getSreamdata);
+          toggleQueueChecker(false);
+        } finally {
+          // Always re-enable the button even if there was an error
+          disabled = false;
+        }
       } else {
-        if (isImageMode) {
-          mediaStreamActions.stop();
+        // Handle stopping - disable button during this process too
+        disabled = true;
+        
+        try {
+          if (isImageMode) {
+            mediaStreamActions.stop();
+          }
+          await lcmLiveActions.stop();
+          toggleQueueChecker(true);
+        } finally {
+          disabled = false;
         }
-        lcmLiveActions.stop();
-        toggleQueueChecker(true);
       }
     } catch (e) {
-      warningMessage = e instanceof Error ? e.message : '';
+      console.error('Error in toggleLcmLive:', e);
+      warningMessage = e instanceof Error ? e.message : 'An unknown error occurred';
       disabled = false;
       toggleQueueChecker(true);
     }
   }
+  
+  // Reconnect function for automatic reconnection
+  async function reconnect() {
+    try {
+      disabled = true;
+      warningMessage = 'Reconnecting...';
+      
+      if (isImageMode) {
+        await mediaStreamActions.stop();
+        await mediaStreamActions.enumerateDevices();
+        await mediaStreamActions.start();
+      }
+      
+      await lcmLiveActions.reconnect(getSreamdata);
+      warningMessage = '';
+      toggleQueueChecker(false);
+    } catch (e) {
+      warningMessage = e instanceof Error ? e.message : 'Reconnection failed';
+      toggleQueueChecker(true);
+    } finally {
+      disabled = false;
+    }
+  }
 </script>
 
 <svelte:head>
@@ -111,6 +162,17 @@
         > and run it on your own GPU.
       </p>
     {/if}
+    
+    {#if $lcmLiveStatus === LCMLiveStatus.ERROR}
+      <p class="text-sm mt-2">
+        <button 
+          class="text-blue-500 underline hover:no-underline" 
+          on:click={reconnect}
+          disabled={disabled}>
+          Try reconnecting
+        </button>
+      </p>
+    {/if}
   </article>
   {#if pipelineParams}
     <article class="my-3 grid grid-cols-1 gap-3 sm:grid-cols-4">
@@ -127,7 +189,9 @@
       </div>
       <div class="sm:col-span-4 sm:row-start-2">
         <Button on:click={toggleLcmLive} {disabled} classList={'text-lg my-1 p-2'}>
-          {#if isLCMRunning}
+          {#if isConnecting}
+            Connecting...
+          {:else if isLCMRunning}
             Stop
           {:else}
             Start

server/config.py

@@ -1,9 +1,10 @@
-from typing import NamedTuple
+from pydantic import BaseModel, field_validator
 import argparse
 import os
+from typing import Annotated
 
 
-class Args(NamedTuple):
+class Args(BaseModel):
     host: str
     port: int
     reload: bool
@@ -13,19 +14,30 @@ class Args(NamedTuple):
     torch_compile: bool
     taesd: bool
     pipeline: str
-    ssl_certfile: str
-    ssl_keyfile: str
+    ssl_certfile: str | None
+    ssl_keyfile: str | None
     sfast: bool
     onediff: bool = False
     compel: bool = False
     debug: bool = False
 
-    def pretty_print(self):
+    def pretty_print(self) -> None:
         print("\n")
-        for field, value in self._asdict().items():
+        for field, value in self.model_dump().items():
             print(f"{field}: {value}")
         print("\n")
 
+    @field_validator("ssl_keyfile")
+    @classmethod
+    def validate_ssl_keyfile(cls, v: str | None, info) -> str | None:
+        """Validate that if ssl_certfile is provided, ssl_keyfile is also provided."""
+        ssl_certfile = info.data.get("ssl_certfile")
+        if ssl_certfile and not v:
+            raise ValueError(
+                "If ssl_certfile is provided, ssl_keyfile must also be provided"
+            )
+        return v
+
 
 MAX_QUEUE_SIZE = int(os.environ.get("MAX_QUEUE_SIZE", 0))
 TIMEOUT = float(os.environ.get("TIMEOUT", 0))
@@ -113,5 +125,5 @@ parser.add_argument(
 )
 parser.set_defaults(taesd=USE_TAESD)
 
-config = Args(**vars(parser.parse_args()))
+config = Args.model_validate(vars(parser.parse_args()))
 config.pretty_print()

server/connection_manager.py

@@ -1,6 +1,7 @@
 from uuid import UUID
 import asyncio
 from fastapi import WebSocket
+from fastapi.websockets import WebSocketDisconnect
 from starlette.websockets import WebSocketState
 import logging
 from typing import Any
@@ -73,44 +74,157 @@ class ConnectionManager:
     def get_user_count(self) -> int:
         return len(self.active_connections)
 
-    def get_websocket(self, user_id: UUID) -> WebSocket:
+    def get_websocket(self, user_id: UUID) -> WebSocket | None:
         user_session = self.active_connections.get(user_id)
         if user_session:
             websocket = user_session["websocket"]
-            if websocket.client_state == WebSocketState.CONNECTED:
+            # Both client_state and application_state should be checked
+            # to ensure the websocket is fully connected and not closing
+            if (websocket.client_state == WebSocketState.CONNECTED and 
+                websocket.application_state == WebSocketState.CONNECTED):
                 return user_session["websocket"]
         return None
 
     async def disconnect(self, user_id: UUID):
-        websocket = self.get_websocket(user_id)
-        if websocket:
-            await websocket.close()
+        # First check if user is in active connections
+        if user_id not in self.active_connections:
+            return
+            
+        # Get the websocket directly from active_connections to avoid get_websocket validation
+        user_session = self.active_connections.get(user_id)
+        if user_session and "websocket" in user_session:
+            websocket = user_session["websocket"]
+            try:
+                # Only attempt close if not already closed
+                if (websocket.client_state != WebSocketState.DISCONNECTED and 
+                    websocket.application_state != WebSocketState.DISCONNECTED):
+                    await websocket.close()
+            except Exception as e:
+                logging.error(f"Error closing websocket for {user_id}: {e}")
+                
+        # Always delete the user to ensure cleanup
         self.delete_user(user_id)
 
     async def send_json(self, user_id: UUID, data: dict):
         try:
             websocket = self.get_websocket(user_id)
             if websocket:
-                await websocket.send_json(data)
+                try:
+                    await websocket.send_json(data)
+                except RuntimeError as e:
+                    error_msg = str(e)
+                    if any(err in error_msg for err in [
+                        "WebSocket is not connected", 
+                        "Cannot call \"send\" once a close message has been sent",
+                        "Cannot call \"receive\" once a close message has been sent",
+                        "WebSocket is disconnected"]):
+                        # The websocket was disconnected or is closing
+                        logging.info(f"WebSocket disconnected/closing for user {user_id}: {error_msg}")
+                        await self.disconnect(user_id)
+                    else:
+                        logging.error(f"Runtime error in send_json: {e}")
+        except WebSocketDisconnect as disconnect_error:
+            # Handle websocket disconnection event
+            code = disconnect_error.code
+            if code == 1006:  # ABNORMAL_CLOSURE
+                logging.info(f"WebSocket abnormally closed for user {user_id} during send: Connection was closed without a proper close handshake")
+            else:
+                logging.info(f"WebSocket disconnected for user {user_id} with code {code} during send: {disconnect_error.reason}")
+            
+            # Always disconnect the user
+            if user_id in self.active_connections:
+                await self.disconnect(user_id)
         except Exception as e:
             logging.error(f"Error: Send json: {e}")
+            # If any send fails, ensure the user gets removed to prevent further errors
+            if user_id in self.active_connections:
+                await self.disconnect(user_id)
 
     async def receive_json(self, user_id: UUID) -> dict | None:
         try:
             websocket = self.get_websocket(user_id)
             if websocket:
-                return await websocket.receive_json()
+                try:
+                    # Receive the raw message and handle JSON parsing manually for better error handling
+                    try:
+                        data = await websocket.receive_json()
+                        # Verify it's a dictionary
+                        if not isinstance(data, dict):
+                            logging.error(f"Expected dict but received {type(data)} from user {user_id}: {data}")
+                            return None
+                        return data
+                    except ValueError as json_err:
+                        # Specific handling for JSON parsing errors
+                        logging.error(f"JSON parsing error for user {user_id}: {json_err}")
+                        return None
+                except RuntimeError as e:
+                    error_msg = str(e)
+                    if any(err in error_msg for err in [
+                        "WebSocket is not connected", 
+                        "Cannot call \"send\" once a close message has been sent",
+                        "Cannot call \"receive\" once a close message has been sent",
+                        "WebSocket is disconnected"]):
+                        # The websocket was disconnected or closing
+                        logging.info(f"WebSocket disconnected/closing for user {user_id}: {error_msg}")
+                        await self.disconnect(user_id)
+                    else:
+                        logging.error(f"Runtime error in receive_json: {e}")
+                    return None
+            return None
+        except WebSocketDisconnect as disconnect_error:
+            # Handle websocket disconnection event (this is a clean, expected path)
+            code = disconnect_error.code
+            if code == 1006:  # ABNORMAL_CLOSURE
+                logging.info(f"WebSocket abnormally closed for user {user_id}: Connection was closed without a proper close handshake")
+            else:
+                logging.info(f"WebSocket disconnected for user {user_id} with code {code}: {disconnect_error.reason}")
+            
+            # Always disconnect the user
+            if user_id in self.active_connections:
+                await self.disconnect(user_id)
             return None
         except Exception as e:
             logging.error(f"Error: Receive json: {e}")
+            # Ensure disconnection on any exception
+            if user_id in self.active_connections:
+                await self.disconnect(user_id)
             return None
 
     async def receive_bytes(self, user_id: UUID) -> bytes | None:
         try:
             websocket = self.get_websocket(user_id)
             if websocket:
-                return await websocket.receive_bytes()
+                try:
+                    return await websocket.receive_bytes()
+                except RuntimeError as e:
+                    error_msg = str(e)
+                    if any(err in error_msg for err in [
+                        "WebSocket is not connected", 
+                        "Cannot call \"send\" once a close message has been sent",
+                        "Cannot call \"receive\" once a close message has been sent",
+                        "WebSocket is disconnected"]):
+                        # The websocket was disconnected or closing
+                        logging.info(f"WebSocket disconnected/closing for user {user_id}: {error_msg}")
+                        await self.disconnect(user_id)
+                    else:
+                        logging.error(f"Runtime error in receive_bytes: {e}")
+                    return None
+            return None
+        except WebSocketDisconnect as disconnect_error:
+            # Handle websocket disconnection event (this is a clean, expected path)
+            code = disconnect_error.code
+            if code == 1006:  # ABNORMAL_CLOSURE
+                logging.info(f"WebSocket abnormally closed for user {user_id}: Connection was closed without a proper close handshake")
+            else:
+                logging.info(f"WebSocket disconnected for user {user_id} with code {code}: {disconnect_error.reason}")
+            
+            # Always disconnect the user
+            if user_id in self.active_connections:
+                await self.disconnect(user_id)
             return None
         except Exception as e:
             logging.error(f"Error: Receive bytes: {e}")
+            # Ensure disconnection on any exception
+            if user_id in self.active_connections:
+                await self.disconnect(user_id)
             return None

server/main.py

@@ -1,4 +1,5 @@
-from fastapi import FastAPI, WebSocket, HTTPException, WebSocketDisconnect
+from fastapi import FastAPI, WebSocket, HTTPException
+from fastapi.websockets import WebSocketDisconnect
 from fastapi.responses import StreamingResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
@@ -9,37 +10,38 @@ from PIL import Image
 import logging
 from config import config, Args
 from connection_manager import ConnectionManager, ServerFullException
-import uuid
 from uuid import UUID
 import time
-from typing import Any, Protocol, runtime_checkable
+from typing import Any, Protocol, runtime_checkable, AsyncGenerator
 from util import pil_to_frame, bytes_to_pil, is_firefox, get_pipeline_class, ParamsModel
 from device import device, torch_dtype
 import asyncio
 import os
 import time
-import torch
-from pydantic import BaseModel, create_model
+
+# Common WebSocket error messages that indicate disconnection
+ERROR_MESSAGES = [
+    "WebSocket is not connected",
+    'Cannot call "send" once a close message has been sent',
+    'Cannot call "receive" once a close message has been sent',
+    "WebSocket is disconnected",
+]
+
 
 @runtime_checkable
 class BasePipeline(Protocol):
     class Info:
         @classmethod
-        def schema(cls) -> dict[str, Any]:
-            ...
+        def schema(cls) -> dict[str, Any]: ...
+
         page_content: str | None
         input_mode: str
-    
+
     class InputParams(ParamsModel):
         @classmethod
-        def schema(cls) -> dict[str, Any]:
-            ...
-        
-        def dict(self) -> dict[str, Any]:
-            ...
-            
-    def predict(self, params: ParamsModel) -> Image.Image | None:
-        ...
+        def schema(cls) -> dict[str, Any]: ...
+
+    def predict(self, params: ParamsModel) -> Image.Image | None: ...
 
 
 THROTTLE = 1.0 / 120
@@ -74,7 +76,22 @@ class App:
                 await handle_websocket_data(user_id)
             except ServerFullException as e:
                 logging.error(f"Server Full: {e}")
+            except WebSocketDisconnect as disconnect_error:
+                # Handle websocket disconnection event
+                code = disconnect_error.code
+                if code == 1006:  # ABNORMAL_CLOSURE
+                    logging.info(f"WebSocket abnormally closed for user {user_id}: Connection was closed without a proper close handshake")
+                else:
+                    logging.info(f"WebSocket disconnected for user {user_id} with code {code}: {disconnect_error.reason}")
+            except RuntimeError as e:
+                if any(err in str(e) for err in ERROR_MESSAGES):
+                    logging.info(f"WebSocket disconnected for user {user_id}: {e}")
+                else:
+                    logging.error(f"Runtime error in websocket endpoint: {e}")
+            except Exception as e:
+                logging.error(f"Unexpected error in websocket endpoint: {e}")
             finally:
+                # Always ensure we disconnect the user
                 await self.conn_manager.disconnect(user_id)
                 logging.info(f"User disconnected: {user_id}")
 
@@ -99,34 +116,73 @@ class App:
                         return
                     data = await self.conn_manager.receive_json(user_id)
                     if data is None:
+                        # Check if the user is still connected - they might have disconnected
+                        if not self.conn_manager.check_user(user_id):
+                            logging.info(
+                                f"User {user_id} disconnected, exiting handle_websocket_data loop"
+                            )
+                            return
+                        continue
+
+                    # Validate that data is a dictionary and has a status field
+                    if not isinstance(data, dict) or "status" not in data:
+                        logging.error(
+                            f"Invalid data format received from user {user_id}: {data}"
+                        )
                         continue
-                        
+
                     if data["status"] == "next_frame":
                         info = self.pipeline.Info()
                         params_data = await self.conn_manager.receive_json(user_id)
                         if params_data is None:
+                            # Check if the user is still connected
+                            if not self.conn_manager.check_user(user_id):
+                                logging.info(
+                                    f"User {user_id} disconnected during params reception"
+                                )
+                                return
                             continue
-                            
+
                         params = self.pipeline.InputParams.model_validate(params_data)
-                        
+
                         if info.input_mode == "image":
                             image_data = await self.conn_manager.receive_bytes(user_id)
-                            if image_data is None or len(image_data) == 0:
+                            if image_data is None:
+                                # Check if the user is still connected
+                                if not self.conn_manager.check_user(user_id):
+                                    logging.info(
+                                        f"User {user_id} disconnected during image reception"
+                                    )
+                                    return
                                 await self.conn_manager.send_json(
                                     user_id, {"status": "send_frame"}
                                 )
                                 continue
-                            
-                            # Create a new Pydantic model with the image field
-                            params_dict = params.model_dump()
-                            params_dict["image"] = bytes_to_pil(image_data)
-                            params = self.pipeline.InputParams.model_validate(params_dict)
+                            if len(image_data) == 0:
+                                await self.conn_manager.send_json(
+                                    user_id, {"status": "send_frame"}
+                                )
+                                continue
+
+                            # Add the image directly to the model using setattr
+                            # This works because we've configured the ParamsModel to allow extra fields
+                            setattr(params, "image", bytes_to_pil(image_data))
 
                         await self.conn_manager.update_data(user_id, params)
                         await self.conn_manager.send_json(user_id, {"status": "wait"})
 
+            except RuntimeError as e:
+                error_msg = str(e)
+                if any(err in error_msg for err in ERROR_MESSAGES):
+                    logging.info(
+                        f"WebSocket disconnected for user {user_id}: {error_msg}"
+                    )
+                else:
+                    logging.error(f"Websocket Runtime Error: {e}, {user_id}")
+                # Ensure disconnect is called
+                await self.conn_manager.disconnect(user_id)
             except Exception as e:
-                logging.error(f"Websocket Error: {e}, {user_id} ")
+                logging.error(f"Websocket Error: {e}, {user_id}")
                 await self.conn_manager.disconnect(user_id)
 
         @self.app.get("/api/queue")
@@ -137,25 +193,58 @@ class App:
         @self.app.get("/api/stream/{user_id}")
         async def stream(user_id: UUID, request: Request) -> StreamingResponse:
             try:
-                async def generate() -> bytes:
+
+                async def generate() -> AsyncGenerator[bytes, None]:
                     last_params: ParamsModel | None = None
                     while True:
+                        # Check if the user is still connected
+                        if not self.conn_manager.check_user(user_id):
+                            logging.info(f"User {user_id} disconnected from stream")
+                            break
+
                         last_time = time.time()
-                        await self.conn_manager.send_json(
-                            user_id, {"status": "send_frame"}
-                        )
+                        try:
+                            await self.conn_manager.send_json(
+                                user_id, {"status": "send_frame"}
+                            )
+                        except Exception as e:
+                            logging.error(f"Error sending to websocket in stream: {e}")
+                            # User might have disconnected
+                            if not self.conn_manager.check_user(user_id):
+                                logging.info(f"User {user_id} disconnected from stream")
+                                break
+                            await asyncio.sleep(THROTTLE)
+                            continue
+
                         params = await self.conn_manager.get_latest_data(user_id)
-                        
-                        if (params is None or 
-                            (last_params is not None and 
-                             params.model_dump() == last_params.model_dump())):
+
+                        if params is None:
                             await asyncio.sleep(THROTTLE)
                             continue
-                            
-                        last_params = params
-                        image = self.pipeline.predict(params)
 
-                        if self.args.safety_checker and self.safety_checker is not None and image is not None:
+                        try:
+                            # Check if the params haven't changed since last time
+                            if (
+                                last_params is not None
+                                and params.model_dump() == last_params.model_dump()
+                            ):
+                                await asyncio.sleep(THROTTLE)
+                                continue
+
+                            last_params = params
+                            image = self.pipeline.predict(params)
+                        except Exception as e:
+                            logging.error(
+                                f"Error processing params for user {user_id}: {e}"
+                            )
+                            await asyncio.sleep(THROTTLE)
+                            continue
+
+                        if (
+                            self.args.safety_checker
+                            and self.safety_checker is not None
+                            and image is not None
+                        ):
                             image, has_nsfw_concept = self.safety_checker(image)
                             if has_nsfw_concept:
                                 image = None
@@ -175,9 +264,34 @@ class App:
                     media_type="multipart/x-mixed-replace;boundary=frame",
                     headers={"Cache-Control": "no-cache"},
                 )
+            except WebSocketDisconnect as disconnect_error:
+                # Handle websocket disconnection event
+                code = disconnect_error.code
+                if code == 1006:  # ABNORMAL_CLOSURE
+                    logging.info(f"WebSocket abnormally closed during streaming for user {user_id}: Connection was closed without a proper close handshake")
+                else:
+                    logging.info(f"WebSocket disconnected during streaming for user {user_id} with code {code}: {disconnect_error.reason}")
+                
+                # Clean disconnection without error response
+                await self.conn_manager.disconnect(user_id)
+                raise HTTPException(status_code=204, detail="Connection closed")
+            except RuntimeError as e:
+                error_msg = str(e)
+                if any(err in error_msg for err in ERROR_MESSAGES):
+                    logging.info(
+                        f"WebSocket disconnected during streaming for user {user_id}: {error_msg}"
+                    )
+                    # Clean disconnection without error response
+                    await self.conn_manager.disconnect(user_id)
+                    raise HTTPException(status_code=204, detail="Connection closed")
+                else:
+                    logging.error(f"Streaming Runtime Error: {e}, {user_id}")
+                    raise HTTPException(status_code=500, detail="Streaming error")
             except Exception as e:
-                logging.error(f"Streaming Error: {e}, {user_id} ")
-                raise HTTPException(status_code=404, detail="User not found")
+                logging.error(f"Streaming Error: {e}, {user_id}")
+                # Always ensure we disconnect the user on error
+                await self.conn_manager.disconnect(user_id)
+                raise HTTPException(status_code=500, detail="Streaming error")
 
         # route to setup frontend
         @self.app.get("/api/settings")
@@ -185,7 +299,7 @@ class App:
             info_schema = self.pipeline.Info.schema()
             info = self.pipeline.Info()
             page_content = ""
-            if hasattr(info, 'page_content') and info.page_content:
+            if hasattr(info, "page_content") and info.page_content:
                 page_content = markdown2.markdown(info.page_content)
 
             input_params = self.pipeline.InputParams.schema()
@@ -234,7 +348,7 @@ if __name__ == "__main__":
     # app = create_app(config)  # Create the app once
 
     uvicorn.run(
-        app,
+        "main:app",
         host=config.host,
         port=config.port,
         reload=config.reload,

server/pipelines/img2img.py

@@ -14,6 +14,7 @@ import psutil
 from config import Args
 from pydantic import BaseModel, Field
 from PIL import Image
+from util import ParamsModel
 import math
 
 base_model = "SimianLuo/LCM_Dreamshaper_v7"
@@ -54,7 +55,7 @@ class Pipeline:
         input_mode: str = "image"
         page_content: str = page_content
 
-    class InputParams(BaseModel):
+    class InputParams(ParamsModel):
         prompt: str = Field(
             default_prompt,
             title="Prompt",

server/pipelines/img2imgFlux.py

@@ -28,6 +28,7 @@ from config import Args
 from pydantic import BaseModel, Field
 from PIL import Image
 from pathlib import Path
+from util import ParamsModel
 import math
 import gc
 
@@ -61,7 +62,7 @@ class Pipeline:
         input_mode: str = "image"
         page_content: str = page_content
 
-    class InputParams(BaseModel):
+    class InputParams(ParamsModel):
         prompt: str = Field(
             default_prompt,
             title="Prompt",

server/pipelines/img2imgSDTurbo.py

@@ -9,10 +9,10 @@ try:
 except:
     pass
 
-import psutil
 from config import Args
 from pydantic import BaseModel, Field
 from PIL import Image
+from util import ParamsModel
 import math
 
 
@@ -55,7 +55,7 @@ class Pipeline:
         input_mode: str = "image"
         page_content: str = page_content
 
-    class InputParams(BaseModel):
+    class InputParams(ParamsModel):
         prompt: str = Field(
             default_prompt,
             title="Prompt",

server/pipelines/img2imgSDXL-Lightning.py

@@ -18,6 +18,7 @@ from huggingface_hub import hf_hub_download
 from config import Args
 from pydantic import BaseModel, Field
 from PIL import Image
+from util import ParamsModel
 import math
 
 base = "stabilityai/stable-diffusion-xl-base-1.0"
@@ -62,7 +63,7 @@ class Pipeline:
         input_mode: str = "image"
         page_content: str = page_content
 
-    class InputParams(BaseModel):
+    class InputParams(ParamsModel):
         prompt: str = Field(
             default_prompt,
             title="Prompt",

server/pipelines/img2imgSDXLTurbo.py

@@ -14,6 +14,7 @@ import psutil
 from config import Args
 from pydantic import BaseModel, Field
 from PIL import Image
+from util import ParamsModel
 import math
 
 base_model = "stabilityai/sdxl-turbo"
@@ -55,7 +56,7 @@ class Pipeline:
         input_mode: str = "image"
         page_content: str = page_content
 
-    class InputParams(BaseModel):
+    class InputParams(ParamsModel):
         prompt: str = Field(
             default_prompt,
             title="Prompt",

server/pipelines/img2imgSDXS512.py

@@ -11,6 +11,7 @@ import psutil
 from config import Args
 from pydantic import BaseModel, Field
 from PIL import Image
+from util import ParamsModel
 import math
 
 base_model = "IDKiro/sdxs-512-0.9"
@@ -51,7 +52,7 @@ class Pipeline:
         input_mode: str = "image"
         page_content: str = page_content
 
-    class InputParams(BaseModel):
+    class InputParams(ParamsModel):
         prompt: str = Field(
             default_prompt,
             title="Prompt",

server/util.py

@@ -1,5 +1,5 @@
 from importlib import import_module
-from typing import Any, TypeVar, type_check_only
+from typing import Any, TypeVar
 from PIL import Image
 import io
 from pydantic import BaseModel
@@ -11,12 +11,17 @@ TPipeline = TypeVar("TPipeline", bound=type[Any])
 
 class ParamsModel(BaseModel):
     """Base model for pipeline parameters."""
-    
+
+    model_config = {
+        "arbitrary_types_allowed": True,
+        "extra": "allow",  # Allow extra attributes for dynamic fields like 'image'
+    }
+
     @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> 'ParamsModel':
+    def from_dict(cls, data: dict[str, Any]) -> "ParamsModel":
         """Create a model instance from dictionary data."""
         return cls.model_validate(data)
-    
+
     def to_dict(self) -> dict[str, Any]:
         """Convert model to dictionary."""
         return self.model_dump()
@@ -25,13 +30,13 @@ class ParamsModel(BaseModel):
 def get_pipeline_class(pipeline_name: str) -> type:
     """
     Dynamically imports and returns the Pipeline class from a specified module.
-    
+
     Args:
         pipeline_name: The name of the pipeline module to import
-        
+
     Returns:
         The Pipeline class from the specified module
-        
+
     Raises:
         ValueError: If the module or Pipeline class isn't found
         TypeError: If Pipeline is not a class