loading indicator spinner update

Cargo.lock

@@ -629,7 +629,6 @@ dependencies = [
  "filesys",
  "hf-hub 0.3.2",
  "image",
- "libloading 0.8.6",
  "log",
  "memory_caching",
  "once_cell",

crates/desktop/src-tauri/Cargo.toml

@@ -48,6 +48,3 @@ tokenizers = "0.15.2"
 [features]
 default = ["cuda"]
 cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
-
-[dev-dependencies]
-libloading = "0.8.6"

crates/desktop/src-tauri/src/main.rs

@@ -1,15 +1,6 @@
 // Prevents additional console window on Windows in release, DO NOT REMOVE!!
 #![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]
 
-use std::fs::File;
-
 fn main() {
-  // NB: This is a test to see the behavior of the application with dynamic loading modules.
-  // We should be able to execute code before loading CUDA and other dependencies.
-  let r = File::create("program_started.txt");
-  match r {
-    Err(err) => println!("Failed to create program_started.txt: {}", err),
-    Ok(_) => {}
-  }
   app_lib::run();
 }

crates/desktop/src-tauri/tests/test_dynamic_library_loading.rs

@@ -1,7 +0,0 @@
-
-#[test]
-fn test_dynamic_library_loading() {
-  let name = libloading::library_filename("cublas");
-
-  assert_eq!(name, "libcublas.so");
-}

frontend/realtime/package-lock.json

@@ -23,7 +23,6 @@
         "react-drag-drop-files": "^2.3.10",
         "react-easy-crop": "^5.1.0",
         "react-router-dom": "^6.26.1",
-        "sonner": "^2.0.1",
         "tailwind-merge": "^2.5.2",
         "uuid": "^10.0.0",
         "yarn": "^1.22.22"
@@ -8290,15 +8289,6 @@
         "url": "https://github.com/sponsors/isaacs"
       }
     },
-    "node_modules/sonner": {
-      "version": "2.0.1",
-      "resolved": "https://registry.npmjs.org/sonner/-/sonner-2.0.1.tgz",
-      "integrity": "sha512-FRBphaehZ5tLdLcQ8g2WOIRE+Y7BCfWi5Zyd8bCvBjiW8TxxAyoWZIxS661Yz6TGPqFQ4VLzOF89WEYhfynSFQ==",
-      "peerDependencies": {
-        "react": "^18.0.0 || ^19.0.0 || ^19.0.0-rc",
-        "react-dom": "^18.0.0 || ^19.0.0 || ^19.0.0-rc"
-      }
-    },
     "node_modules/source-map-js": {
       "version": "1.2.1",
       "resolved": "https://registry.npmjs.org/source-map-js/-/source-map-js-1.2.1.tgz",

frontend/realtime/package.json

@@ -26,7 +26,6 @@
     "react-drag-drop-files": "^2.3.10",
     "react-easy-crop": "^5.1.0",
     "react-router-dom": "^6.26.1",
-    "sonner": "^2.0.1",
     "tailwind-merge": "^2.5.2",
     "uuid": "^10.0.0",
     "yarn": "^1.22.22"

frontend/realtime/paint_py_server/paint_test.py

@@ -1,272 +0,0 @@
-import asyncio
-import base64
-import json
-import logging
-import pathlib
-import time
-from io import BytesIO
-from typing import Dict, Any, Optional, Tuple, Callable, List
-from dataclasses import dataclass
-from dataclasses_json import dataclass_json
-from enum import Enum, auto
-import base64
-from PIL import Image
-import io
-import torch
-from diffusers import AutoPipelineForImage2Image, LCMScheduler
-from diffusers.utils import make_image_grid, load_image
-from sd_embed.embedding_funcs import get_weighted_text_embeddings_sdxl
-import torch
-import websockets
-from PIL import Image
-from diffusers import (
-    StableDiffusionXLPipeline,
-    StableDiffusionXLImg2ImgPipeline,
-    AutoPipelineForImage2Image,
-    LCMScheduler,
-    UNet2DConditionModel,
-)
-import torch
-import torchvision.transforms as transforms
-import gc
-
-transform = transforms.Compose([
-    transforms.ToTensor(),  # This converts PIL Image to tensor and scales to [0, 1]
-])
-
-@dataclass_json
-@dataclass
-class ProgressMessage:
-    message: str
-    progress: float
-    type: str = "progress"
-    error: Optional[str] = None
-
-@dataclass_json
-@dataclass
-class LoadModelRequest:
-    model_path: str
-    lora_path: Optional[str] = None
-    command: str = "load_model"
-
-@dataclass_json
-@dataclass
-class GenerateRequest:
-    image: str
-    prompt: str
-    command: str = "generate"
-    strength: float = 0.6
-    guidance_scale: float = 2.0
-    num_inference_steps: int = 4
-    height: int = 1024
-    width: int = 1024
-    lora_strength: float = 1.0
-
-@dataclass_json
-@dataclass
-class GenerateResponse:
-    image: str
-    type: str = "result"
-
-def save_image(image: Image.Image, base_path: str = "outputs") -> str:
-    """Save a PIL Image to disk with a timestamp-based filename."""
-    try:
-        # Create outputs directory if it doesn't exist
-        output_dir = pathlib.Path(base_path)
-        output_dir.mkdir(parents=True, exist_ok=True)
-        
-        # Generate filename with timestamp
-        timestamp = time.strftime("%Y%m%d-%H%M%S")
-        filename = f"generated-{timestamp}.png"
-        filepath = output_dir / filename
-        
-        # Save the image
-        print(f"Saving image to: {filepath}")
-        image.save(str(filepath), format="PNG")
-        print(f"Successfully saved image to: {filepath}")
-        return str(filepath)
-    except Exception as e:
-        print(f"Error saving image: {e}")
-        return ""
-
-
-async def send_progress(websocket, message: str, progress: float = 0.0, error: str = None):
-    """Send progress updates to the client."""
-    progress_msg = ProgressMessage(
-        message=message,
-        progress=progress,
-        error=error
-    )
-    await websocket.send(progress_msg.to_json())
-
-async def load_model(model_path: str, lora_path: str = None) -> StableDiffusionXLImg2ImgPipeline:
-    """Load the model with progress reporting."""
-    try:
-        # Load base SDXL pipeline
-        pipe = StableDiffusionXLPipeline.from_single_file(
-            model_path,
-            torch_dtype=torch.float16,
-            use_safetensors=True
-        ).to("cuda")
-        
-        # Convert to img2img pipeline
-        pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
-        
-        # Set scheduler
-        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-        
-        # Load LCM-LoRA weights if provided
-        if lora_path:
-            pipe.load_lora_weights(
-                "F:/ComfyUI_windows_portable_nvidia/ComfyUI_windows_portable/ComfyUI/models/loras/LCM_LoRA_Weights_SDXL.safetensors",
-                adapter_name="lcm"
-            )
-            pipe.load_lora_weights(lora_path, adapter_name="lora")
-            pipe.set_adapters(["lcm", "lora"], adapter_weights=[1.0, 1.0])
-            #pipe.set_adapters(["lcm"], adapter_weights=[1.0])
-        #pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-        return pipe
-    except Exception as e:
-        raise Exception(f"Error loading model: {str(e)}")
-
-#https://huggingface.co/blog/simple_sdxl_optimizations
-async def generate_image(
-    pipe: StableDiffusionXLImg2ImgPipeline,
-    request: GenerateRequest
-) -> str:
-    """Generate image and return base64 string."""
-    try:
-        # Clear CUDA cache before processing
-        torch.cuda.empty_cache()
-        
-        image_stream = io.BytesIO(base64.b64decode(request.image))
-      
-        init_image = Image.open(image_stream).convert('RGB')
-        init_image = init_image.resize((request.width, request.height))  # width, height
-        init_image = transform(init_image).to('cuda')
-        init_image = init_image.unsqueeze(0)
-
-        (
-            prompt_embeds
-            , prompt_neg_embeds
-            , pooled_prompt_embeds
-            , negative_pooled_prompt_embeds
-        ) = get_weighted_text_embeddings_sdxl(
-            pipe, 
-            prompt = request.prompt, 
-            neg_prompt = "bad quality, poorly Rendered face, poorly drawn face, poorly facial details, poorly drawn hands, poorly rendered hands, low resolution, Images cut out at the top, left, right, bottom, bad composition, mutated body parts, blurry image, disfigured, oversaturated, bad anatomy, deformed body features"
-        )
-        
-        # Generate image
-        with torch.inference_mode():
-            output_image = pipe(
-                prompt_embeds = prompt_embeds, 
-                negative_prompt_embeds = prompt_neg_embeds, 
-                pooled_prompt_embeds = pooled_prompt_embeds, 
-                negative_pooled_prompt_embeds = negative_pooled_prompt_embeds,
-                image=init_image,
-                num_inference_steps=request.num_inference_steps,
-                guidance_scale=request.guidance_scale,
-                strength=request.strength,
-                generator=torch.manual_seed(0),
-                height=request.height, width=request.width,
-            ).images[0]
-
-        # Clean up CUDA tensors
-        del init_image, prompt_embeds, prompt_neg_embeds,pooled_prompt_embeds, negative_pooled_prompt_embeds
-        torch.cuda.empty_cache()
-        gc.collect()
-
-        # Convert to base64
-        buffered = BytesIO()
-        output_image.save(buffered, format="PNG")
-        buffered.seek(0)
-        img_str = base64.b64encode(buffered.getvalue()).decode()
-        
-        return f"data:image/png;base64,{img_str}"
-    except Exception as e:
-        raise Exception(f"Error generating image: {str(e)} at line {e.__traceback__.tb_lineno}")
-
-async def handle_client(websocket):
-    """Handle WebSocket client connection."""
-    pipe = None
-    try:
-        async for message in websocket:
-            try:
-                # Clear memory before processing new request
-                torch.cuda.empty_cache()
-                gc.collect()
-
-                data = json.loads(message)
-                command = data.get("command")
-
-                if command == "load_model":
-                    await send_progress(websocket, "Loading model...", 0.0)
-                    request = LoadModelRequest.from_json(message)
-                    
-                    if pipe is not None:
-                        # Clean up existing model
-                        del pipe
-                        torch.cuda.empty_cache()
-                        gc.collect()
-                    
-                    pipe = await load_model(request.model_path, request.lora_path)
-                    await send_progress(websocket, "Model loaded successfully", 1.0)
-
-                elif command == "generate":
-                    if not pipe:
-                        raise Exception("Model not loaded. Please load model first.")
-                    
-                    await send_progress(websocket, "Generating image...", 0.0)
-                    request = GenerateRequest.from_json(message)
-                    
-                    # Use a timeout for generation to prevent hanging
-                    try:
-                        result = await asyncio.wait_for(
-                            generate_image(pipe=pipe, request=request),
-                            timeout=10.0  # 30 second timeout
-                        )
-                        response = GenerateResponse(image=result)
-                        await websocket.send(response.to_json())
-                    except asyncio.TimeoutError:
-                        # Force cleanup if generation times out
-                        torch.cuda.empty_cache()
-                        gc.collect()
-                        raise Exception("Generation timed out")
-                    
-                else:
-                    raise Exception(f"Unknown command: {command}")
-                
-            except Exception as e:
-                await send_progress(websocket, "Error processing request", 0.0, str(e))
-                # Clean up after error
-                torch.cuda.empty_cache()
-                gc.collect()
-                
-    except websockets.exceptions.ConnectionClosed:
-        print("Client connection closed")
-    except Exception as e:
-        await send_progress(websocket, "Fatal error", 0.0, str(e))
-    finally:
-        # Clean up resources when client disconnects
-        if pipe is not None:
-            del pipe
-        torch.cuda.empty_cache()
-        gc.collect()
-
-async def server():
-    """Start WebSocket server."""
-    async with websockets.serve(
-        handle_client,
-        "localhost",
-        8765,
-        max_size=None,
-        ping_interval=None,  # Disable ping/pong to prevent timeouts
-        max_queue=10  # Limit pending connections
-    ):
-        print("Server started on ws://localhost:8765")
-        await asyncio.Future()  # run forever
-
-if __name__ == "__main__":
-    # test()  # Comment out test function
-    asyncio.run(server())

frontend/realtime/paint_py_server/server.py

@@ -0,0 +1,527 @@
+import asyncio
+import base64
+import json
+import logging
+import pathlib
+import time
+from io import BytesIO
+from typing import Dict, Any, Optional, Tuple, Callable, List
+
+import torch
+import websockets
+from PIL import Image
+from diffusers import (
+    StableDiffusionXLPipeline,
+    StableDiffusionXLImg2ImgPipeline,
+    AutoPipelineForImage2Image,
+    LCMScheduler,
+    UNet2DConditionModel,
+)
+from tqdm import tqdm
+from huggingface_hub import HfApi, hf_hub_download
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# Global variables
+connected = False
+pipe = None
+model_initialized = False
+current_settings = {}
+
+# Custom TQDMProgressCallback for tracking download progress
+class ProgressCallback:
+    def __init__(self, websocket):
+        self.websocket = websocket
+        self.total = 0
+        self.current = 0
+        self.last_update_time = 0
+        self.last_progress = 0
+        self.stage = "downloading"
+        self.current_file = ""
+    
+    async def __call__(self, current, total, file_name=""):
+        self.total = total
+        self.current = current
+        self.current_file = file_name or self.current_file
+        
+        # Only send updates at most once every 500ms to avoid flooding the connection
+        current_time = time.time()
+        progress = int(100 * current / total) if total > 0 else 0
+        
+        if (current_time - self.last_update_time > 0.5 and progress != self.last_progress) or current >= total:
+            self.last_update_time = current_time
+            self.last_progress = progress
+            
+            await self.websocket.send(json.dumps({
+                'type': 'progress_update',
+                'stage': self.stage,
+                'progress': progress,
+                'file': self.current_file
+            }))
+
+# Pure functions
+def decode_base64_to_image(base64_string: str) -> Image.Image:
+    """Decode a base64 string to a PIL Image."""
+    try:
+        image_data = base64.b64decode(base64_string)
+        return Image.open(BytesIO(image_data))
+    except Exception as e:
+        logger.error(f"Error decoding base64 image: {e}")
+        raise ValueError(f"Invalid base64 image: {e}")
+
+def encode_image_to_base64(image: Image.Image) -> str:
+    """Encode a PIL Image to a base64 string."""
+    try:
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")
+        return base64.b64encode(buffered.getvalue()).decode('utf-8')
+    except Exception as e:
+        logger.error(f"Error encoding image to base64: {e}")
+        raise ValueError(f"Error encoding image: {e}")
+
+async def load_model(sdxl_checkpoint_path: pathlib.Path, progress_callback: Optional[ProgressCallback] = None) -> Tuple[AutoPipelineForImage2Image, bool]:
+    """Load the SDXL model with LCM scheduler and report progress."""
+    try:
+        logger.info(f"Loading model from {sdxl_checkpoint_path}")
+        
+        if progress_callback:
+            progress_callback.stage = "loading_unet"
+            progress_callback.current_file = "latent-consistency/lcm-sdxl"
+            await progress_callback(0, 100)
+        
+        # Load the UNet with LCM
+        unet = UNet2DConditionModel.from_pretrained(
+            "latent-consistency/lcm-sdxl",
+            torch_dtype=torch.float16,
+            variant="fp16",
+            use_safetensors=True,
+        )
+        
+        if progress_callback:
+            progress_callback.stage = "loading_base_model"
+            progress_callback.current_file = str(sdxl_checkpoint_path)
+            await progress_callback(25, 100)
+        
+        # Load the base model using the checkpoint path
+        pipe = AutoPipelineForImage2Image.from_pretrained(
+            "stabilityai/stable-diffusion-xl-base-1.0",
+            unet=unet,
+            torch_dtype=torch.float16,
+            variant="fp16",
+           use_safetensors=True
+        ).to("cuda")
+        
+        if progress_callback:
+            progress_callback.stage = "configuring_scheduler"
+            await progress_callback(75, 100)
+        
+        # Replace the scheduler with LCM scheduler
+        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+        
+        if progress_callback:
+            progress_callback.stage = "loading_lcm_adapter"
+            progress_callback.current_file = "latent-consistency/lcm-lora-sdxl"
+            await progress_callback(85, 100)
+        
+        # Load the default LCM adapter
+        pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl", adapter_name="lcm",from_safetensors=True)
+        
+        if progress_callback:
+            await progress_callback(100, 100)
+        
+        return pipe, True
+    except Exception as e:
+        logger.error(f"Error loading model: {e}")
+        if progress_callback:
+            progress_callback.stage = "error"
+            progress_callback.current_file = str(e)
+            await progress_callback(0, 100)
+        return None, False
+
+async def load_lora(pipe: AutoPipelineForImage2Image, lora_path: pathlib.Path, adapter_name: str = "lora", progress_callback: Optional[ProgressCallback] = None) -> bool:
+    """Load a LoRA adapter into the pipeline and report progress."""
+    try:
+        logger.info(f"Loading LoRA from {lora_path} as {adapter_name}")
+        
+        if progress_callback:
+            progress_callback.stage = "loading_lora"
+            progress_callback.current_file = str(lora_path)
+            await progress_callback(0, 100)
+        
+        # Load the LoRA weights
+        pipe.load_lora_weights(lora_path, adapter_name=adapter_name)
+        
+        if progress_callback:
+            await progress_callback(100, 100)
+            
+        return True
+    except Exception as e:
+        logger.error(f"Error loading LoRA: {e}")
+        if progress_callback:
+            progress_callback.stage = "error"
+            progress_callback.current_file = str(e)
+            await progress_callback(0, 100)
+        return False
+
+def generate_image(
+    pipe: AutoPipelineForImage2Image,
+    image: Image.Image,
+    prompt: str,
+    lora_strength: float = 1.0,
+    image_to_image_strength: float = 1.0,
+    width: int = 1024,
+    height: int = 1024,
+    num_inference_steps: int = 4,
+    guidance_scale: float = 1.0,
+    seed: Optional[int] = None
+) -> Optional[Image.Image]:
+    """Generate an image using the LCM model."""
+    try:
+        # Resize input image to match desired dimensions
+        image = image.resize((width, height))
+        
+        # Set up generator for reproducibility if seed is provided
+        generator = torch.manual_seed(seed) if seed is not None else None
+        
+        # Set lora adapter weights if available
+        if hasattr(pipe, 'set_adapters') and len(pipe.adapters) > 0:
+            # If we have both lcm and lora adapters
+            if "lcm" in pipe.adapters and "lora" in pipe.adapters:
+                pipe.set_adapters(["lcm", "lora"], adapter_weights=[1.0, lora_strength])
+            # If we only have lcm adapter
+            elif "lcm" in pipe.adapters:
+                pipe.set_adapters(["lcm"], adapter_weights=[1.0])
+        
+        # Generate the image
+        result = pipe(
+            prompt=prompt,
+            image=image,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            strength=image_to_image_strength,
+            width=width,
+            height=height,
+            generator=generator
+        )
+        
+        return result.images[0]
+    except Exception as e:
+        logger.error(f"Error generating image: {e}")
+        return None
+
+async def handle_client(websocket):
+    """Handle a client connection."""
+    global connected, pipe, model_initialized, current_settings
+    
+    if connected:
+        await websocket.send(json.dumps({
+            'type': 'connection_error',
+            'error': "Another client is already connected"
+        }))
+        await websocket.close(1008, "Another client is already connected")
+        return
+    
+    connected = True
+    logger.info("Client connected")
+    
+    # Create a progress callback for this connection
+    progress_callback = ProgressCallback(websocket)
+    
+    try:
+        async for message in websocket:
+            try:
+                data = json.loads(message)
+                logger.info(f"Received message type: {data.get('type', 'unknown')}")
+                logger.info(f"Full payload: {json.dumps(data, indent=2)}")
+                
+                # Check message type
+                message_type = data.get('type', '')
+                
+                # Handle setup request - initializes the model and settings
+                if message_type == 'setup':
+                    logger.info("Processing setup message")
+                    model_data = data.get('model', {})
+                    lora_data = data.get('lora')
+                    device = data.get('device', 'cuda')
+                    
+                    logger.info(f"Model data: {json.dumps(model_data, indent=2)}")
+                    logger.info(f"LoRA data: {json.dumps(lora_data, indent=2) if lora_data else 'None'}")
+                    logger.info(f"Device: {device}")
+                    
+                    # Extract setup parameters with detailed logging
+                    model_path = model_data.get('path', '')
+                    logger.info(f"Raw model path from message: '{model_path}'")
+                    
+                    # Check if the path is valid
+                    if not model_path:
+                        logger.error("Model path is empty")
+                        await websocket.send(json.dumps({
+                            'type': 'setup_response',
+                            'success': False,
+                            'error': 'Model path is required'
+                        }))
+                        continue
+                    
+                    # Convert to Path object
+                    sdxl_checkpoint_path = pathlib.Path(model_path)
+                    logger.info(f"Converted to pathlib.Path: '{sdxl_checkpoint_path}'")
+                    logger.info(f"Path is absolute: {sdxl_checkpoint_path.is_absolute()}")
+                    logger.info(f"Path exists check: {sdxl_checkpoint_path.exists()}")
+                    
+                    # If the path exists as a file
+                    if sdxl_checkpoint_path.is_file():
+                        logger.info(f"Path is a file: Yes")
+                        logger.info(f"File size: {sdxl_checkpoint_path.stat().st_size / (1024*1024):.2f} MB")
+                    else:
+                        logger.info(f"Path is a file: No")
+                    
+                    # Store current settings
+                    current_settings = {
+                        'sdxl_checkpoint_path': str(sdxl_checkpoint_path),
+                        'lora_path': str(lora_data['path']) if lora_data else None,
+                        'lora_alpha': float(lora_data['alpha']) if lora_data else 1.0,
+                        'device': device
+                    }
+                    logger.info(f"Current settings: {json.dumps(current_settings, indent=2)}")
+                    
+                    # Validate required parameters
+                    if not sdxl_checkpoint_path.exists():
+                        logger.error(f"Path does not exist: {sdxl_checkpoint_path}")
+                        # Try alternative path format
+                        alt_path = pathlib.Path(model_path.replace('\\', '/'))
+                        logger.info(f"Trying alternative path format: {alt_path}")
+                        logger.info(f"Alternative path exists check: {alt_path.exists()}")
+                        
+                        if alt_path.exists():
+                            logger.info("Using alternative path format")
+                            sdxl_checkpoint_path = alt_path
+                        else:
+                            await websocket.send(json.dumps({
+                                'type': 'setup_response',
+                                'success': False,
+                                'error': f'SDXL checkpoint path does not exist: {sdxl_checkpoint_path}'
+                            }))
+                            continue
+                    
+                    # Load the model with progress reporting
+                    logger.info(f"Loading model from: {sdxl_checkpoint_path}")
+                    pipe, success = await load_model(sdxl_checkpoint_path, progress_callback)
+                    
+                    if not success:
+                        await websocket.send(json.dumps({
+                            'type': 'setup_response',
+                            'success': False,
+                            'error': 'Failed to load model'
+                        }))
+                        continue
+                    
+                    # Load LoRA if specified
+                    lora_success = True
+                    if lora_data and lora_data['path'] and pathlib.Path(lora_data['path']).exists():
+                        lora_success = await load_lora(pipe, pathlib.Path(lora_data['path']), progress_callback=progress_callback)
+                        if not lora_success:
+                            await websocket.send(json.dumps({
+                                'type': 'setup_response',
+                                'success': False,
+                                'error': f'Failed to load LoRA from {pathlib.Path(lora_data["path"])}'
+                            }))
+                            continue
+                    
+                    model_initialized = True
+                    await websocket.send(json.dumps({
+                        'type': 'setup_response',
+                        'success': True,
+                        'message': 'Model initialized successfully',
+                        'settings': current_settings
+                    }))
+                
+                # Handle settings update - changes settings without full reinitialization
+                elif message_type == 'update_settings':
+                    if not model_initialized:
+                        await websocket.send(json.dumps({
+                            'type': 'update_settings_response',
+                            'success': False,
+                            'error': 'Model not initialized. Send a setup request first.'
+                        }))
+                        continue
+                    
+                    settings_data = data.get('settings', {})
+                    updates_applied = []
+                    errors = []
+
+                    # Check for LoRA path update
+                    if 'lora_path' in settings_data and settings_data['lora_path'] != current_settings.get('lora_path'):
+                        new_lora_path = pathlib.Path(settings_data['lora_path']) if settings_data['lora_path'] else None
+                        
+                        if new_lora_path and new_lora_path.exists():
+                            # Load the new LoRA with progress reporting
+                            lora_success = await load_lora(pipe, new_lora_path, progress_callback=progress_callback)
+                            if lora_success:
+                                current_settings['lora_path'] = str(new_lora_path)
+                                updates_applied.append('lora_path')
+                            else:
+                                errors.append(f'Failed to load LoRA from {new_lora_path}')
+                        elif new_lora_path:
+                            errors.append(f'LoRA path does not exist: {new_lora_path}')
+                        else:
+                            # If LoRA path is None, we'll just update the setting
+                            current_settings['lora_path'] = None
+                            updates_applied.append('lora_path')
+                    
+                    # Handle any other settings that can be updated without reloading the model
+                    # For example, default generation parameters
+                    for param in ['default_lora_strength', 'default_image_to_image_strength',
+                                 'default_width', 'default_height', 'default_steps',
+                                 'default_guidance_scale']:
+                        if param in settings_data:
+                            current_settings[param] = settings_data[param]
+                            updates_applied.append(param)
+                    
+                    # Send response
+                    if errors:
+                        await websocket.send(json.dumps({
+                            'type': 'update_settings_response',
+                            'success': False,
+                            'error': ', '.join(errors),
+                            'updates_applied': updates_applied,
+                            'current_settings': current_settings
+                        }))
+                    else:
+                        await websocket.send(json.dumps({
+                            'type': 'update_settings_response',
+                            'success': True,
+                            'message': 'Settings updated successfully',
+                            'updates_applied': updates_applied,
+                            'current_settings': current_settings
+                        }))
+                
+                # Handle generation request
+                elif message_type == 'generate':
+                    if not model_initialized:
+                        await websocket.send(json.dumps({
+                            'type': 'generate_response',
+                            'success': False,
+                            'error': 'Model not initialized. Send a setup request first.'
+                        }))
+                        continue
+                    
+                    generate_data = data.get('generate', {})
+                    
+                    # Extract generation parameters
+                    base64_image = generate_data.get('image')
+                    prompt = generate_data.get('prompt', '')
+                    
+                    # Use default settings if not provided
+                    lora_strength = float(generate_data.get('lora_strength', current_settings.get('default_lora_strength', 1.0)))
+                    image_to_image_strength = float(generate_data.get('image_to_image_strength', current_settings.get('default_image_to_image_strength', 1.0)))
+                    width = int(generate_data.get('generated_image_width', current_settings.get('default_width', 1024)))
+                    height = int(generate_data.get('generated_image_height', current_settings.get('default_height', 1024)))
+                    num_inference_steps = int(generate_data.get('num_inference_steps', current_settings.get('default_steps', 4)))
+                    guidance_scale = float(generate_data.get('guidance_scale', current_settings.get('default_guidance_scale', 1.0)))
+                    seed = int(generate_data.get('seed')) if 'seed' in generate_data else None
+                    
+                    # Validate required parameters
+                    if not base64_image:
+                        await websocket.send(json.dumps({
+                            'type': 'generate_response',
+                            'success': False,
+                            'error': 'Missing base64 image'
+                        }))
+                        continue
+                    
+                    # Decode the input image
+                    input_image = decode_base64_to_image(base64_image)
+                    
+                    # Update for generation in progress
+                    await websocket.send(json.dumps({
+                        'type': 'progress_update',
+                        'stage': 'generating',
+                        'progress': 0,
+                        'file': ""
+                    }))
+                    
+                    # Generate the image
+                    output_image = generate_image(
+                        pipe=pipe,
+                        image=input_image,
+                        prompt=prompt,
+                        lora_strength=lora_strength,
+                        image_to_image_strength=image_to_image_strength,
+                        width=width,
+                        height=height,
+                        num_inference_steps=num_inference_steps,
+                        guidance_scale=guidance_scale,
+                        seed=seed
+                    )
+                    
+                    # Update for generation complete
+                    await websocket.send(json.dumps({
+                        'type': 'progress_update',
+                        'stage': 'generating',
+                        'progress': 100,
+                        'file': ""
+                    }))
+                    
+                    if output_image is None:
+                        await websocket.send(json.dumps({
+                            'type': 'generate_response',
+                            'success': False,
+                            'error': 'Failed to generate image'
+                        }))
+                        continue
+                    
+                    # Encode the output image to base64
+                    output_base64 = encode_image_to_base64(output_image)
+                    
+                    # Send the result back to the client
+                    await websocket.send(json.dumps({
+                        'type': 'generate_response',
+                        'success': True,
+                        'image': output_base64
+                    }))
+                
+                # Handle status request
+                elif message_type == 'status':
+                    await websocket.send(json.dumps({
+                        'type': 'status_response',
+                        'success': True,
+                        'model_initialized': model_initialized,
+                        'current_settings': current_settings
+                    }))
+                
+                else:
+                    await websocket.send(json.dumps({
+                        'type': 'error',
+                        'error': f'Unknown message type: {message_type}'
+                    }))
+                    
+            except json.JSONDecodeError:
+                await websocket.send(json.dumps({
+                    'type': 'error',
+                    'error': 'Invalid JSON'
+                }))
+            except Exception as e:
+                logger.error(f"Error processing message: {e}")
+                await websocket.send(json.dumps({
+                    'type': 'error',
+                    'error': str(e)
+                }))
+    except websockets.exceptions.ConnectionClosed:
+        logger.info("Connection closed")
+    finally:
+        connected = False
+        logger.info("Client disconnected")
+
+async def main():
+    """Start the WebSocket server."""
+    server = await websockets.serve(handle_client, "localhost", 8765)
+    logger.info("Server started on ws://localhost:8765")
+    await server.wait_closed()
+
+if __name__ == "__main__":
+    asyncio.run(main()) 

frontend/realtime/paint_py_server/test_copy.py

@@ -1,101 +0,0 @@
-import asyncio
-import base64
-import json
-import logging
-import pathlib
-import time
-from io import BytesIO
-from typing import Dict, Any, Optional, Tuple, Callable, List
-
-import torch
-from diffusers import AutoPipelineForImage2Image, LCMScheduler
-from diffusers.utils import make_image_grid, load_image
-
-import torch
-import websockets
-from PIL import Image
-from diffusers import (
-    StableDiffusionXLPipeline,
-    StableDiffusionXLImg2ImgPipeline,
-    AutoPipelineForImage2Image,
-    LCMScheduler,
-    UNet2DConditionModel,
-)
-
-def save_image(image: Image.Image, base_path: str = "outputs") -> str:
-    """Save a PIL Image to disk with a timestamp-based filename."""
-    try:
-        # Create outputs directory if it doesn't exist
-        output_dir = pathlib.Path(base_path)
-        output_dir.mkdir(parents=True, exist_ok=True)
-        
-        # Generate filename with timestamp
-        timestamp = time.strftime("%Y%m%d-%H%M%S")
-        filename = f"generated-{timestamp}.png"
-        filepath = output_dir / filename
-        
-        # Save the image
-        print(f"Saving image to: {filepath}")
-        image.save(str(filepath), format="PNG")
-        print(f"Successfully saved image to: {filepath}")
-        return str(filepath)
-    except Exception as e:
-        print(f"Error saving image: {e}")
-        return ""
-
-def test():
-    # Load base SDXL pipeline from .safetensors file
-    pipe = StableDiffusionXLPipeline.from_single_file(
-        f"C:/Users/Tensor/Downloads/animagineXL40_v4Opt.safetensors",
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-    ).to("cuda")
-
-    # Convert to img2img pipeline
-    pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
-
-    # Set scheduler
-    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-
-    # Load LCM-LoRA weights
-    pipe.load_lora_weights(
-        "F:/ComfyUI_windows_portable_nvidia/ComfyUI_windows_portable/ComfyUI/models/loras/LCM_LoRA_Weights_SDXL.safetensors",
-        adapter_name="lcm"
-    )
-    pipe.set_adapters(["lcm"], adapter_weights=[1.0])
-
-    # prepare image
-    url = "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/7462a0ac-637d-467b-baa9-34c5f54696b4/original=true,quality=90/60277788.jpeg"
-    init_image = load_image(url)
-    prompt = "safe_pos, 1girl, solo, solo focus, 8k, masterpiece, hires, absurdres, splash art, gradient, looking at viewer, medium breasts, vibrant yellow hair, purple hair tips, (shoulder freckles:1.2), nude, athletic, confident, hand on waist, extremely long hair, dancing, delicate gold jewelry, thin gold bangles, belly button chain, action shot,, head tilt, hair censor, spotlight, tottotonero"
-
-    # pass prompt and image to pipeline
-    generator = torch.manual_seed(0)
-    output_image = pipe(
-        prompt,
-        image=init_image,
-        num_inference_steps=4,
-        guidance_scale=2,
-        strength=0.6,
-        generator=generator
-    ).images[0]
-
-    # Save both input and output images
-    save_image(init_image, "outputs/input")
-    output_path = save_image(output_image, "outputs/output")
-    print(f"Output image saved to: {output_path}")
-
-    # Create and save the grid
-    grid = make_image_grid([init_image, output_image], rows=1, cols=2)
-    grid_path = save_image(grid, "outputs/grid")
-    print(f"Grid image saved to: {grid_path}")
-
-async def start_server():
-    async with websockets.serve(
-        "localhost", 
-        8765 # Increase to 10MB
-    ):
-        await asyncio.Future()  # run forever
-
-if __name__ == "__main__":
-    test()

frontend/realtime/src/KonvaApp/Engine.ts

@@ -30,7 +30,8 @@ import { VideoExtractionHandler } from "./EngineUtitlities/VideoExtractionHandle
 import { RealTimeDrawEngine } from "./RenderingPrimitives/RealTimeDrawEngine";
 import { NodeColor } from "~/signals/uiEvents/toolbarNode";
 import { ShapeType } from "./Nodes";
-
+// for testing loading files from system
+// import { FileUtilities } from "./FileUtilities/FileUtilities";
 
 export interface RenderingOptions {
   artstyle: string;
@@ -63,7 +64,6 @@ export class Engine {
   private selectorSquare: SelectorSquare;
   private loadingVideosProvider: LoadingVideosProvider;
   private matteBox: MatteBox;
-  private backgroundNode: ShapeNode | null = null;
 
   private sceneManager: SceneManager;
   private undoStackManager: UndoStackManager;
@@ -196,20 +196,19 @@ export class Engine {
         this.selectorSquare.disable(); // this breaks the paint if you use it in the wrong spot
         this.selectionManager.disable(); // this prevents selection.
         this.realTimeDrawEngine.enablePaintMode();
+        this.realTimeDrawEngine.disableDragging();
 
         uiAccess.toolbarMain.enable();
         uiAccess.toolbarMain.changeButtonState(ToolbarMainButtonNames.SELECT, {
           active: false,
         });
-        uiAccess.toolbarMain.changeButtonState(ToolbarMainButtonNames.PAINT, {
-          active: true,
-        });
         this.matteBox.disable();
         return;
       }
       case AppModes.SELECT: {
         console.log("APPMODE: SELECT");
         this.realTimeDrawEngine.disablePaintMode();
+        this.realTimeDrawEngine.enableDragging();
 
         this.selectorSquare.enable();
         this.selectionManager.enable();
@@ -217,9 +216,6 @@ export class Engine {
         uiAccess.toolbarMain.changeButtonState(ToolbarMainButtonNames.SELECT, {
           active: true,
         });
-        uiAccess.toolbarMain.changeButtonState(ToolbarMainButtonNames.PAINT, {
-          active: false,
-        });
         this.matteBox.disable();
         document.body.style.cursor = "default";
         return;
@@ -314,9 +310,66 @@ export class Engine {
         });
       }
     });
-   
+    uiEvents.toolbarNode.CRHOMA.onClick(() => {
+      const nodes = this.selectionManager.getSelectedNodes();
+      if (nodes.size > 1) {
+        uiAccess.dialogError.show({
+          title: "Error: Background Removal",
+          message:
+            "Please do not select more than 1 item for the Background Removal feature, we can only apply Background Removal to 1 item at a time",
+        });
+        return;
+      }
+      const node = nodes.values().next().value;
+      try {
+        if (node instanceof VideoNode) {
+          const nodeChromaProps = node.getChroma();
+          uiAccess.dialogChromakey.show({
+            isChromakeyEnabled: nodeChromaProps.isChromakeyEnabled,
+            chromakeyColor: nodeChromaProps.chromakeyColor,
+          });
+        } else {
+          throw new Error();
+        }
+      } catch {
+        uiAccess.dialogError.show({
+          title: "Error: Background Removal",
+          message: "This item is not compatible is Background Removal",
+        });
+      }
+    });
 
-  
+    uiEvents.toolbarNode.SEGMENTATION.onClick(async () => {
+      if (this.segmentationButtonCanBePressed == false) {
+        console.log("VideoExtraction Button DEBOUNCED ");
+        return;
+      }
+      console.log("VideoExtraction Button Clicked ACCEPTED");
+
+      // Gating for an appropriate selection
+      const nodes = this.selectionManager.getSelectedNodes();
+      if (nodes.size > 1) {
+        // display error that segmentation cannot be done on more than 1 at a time.
+        uiAccess.dialogError.show({
+          title: "Error: Video Extraction",
+          message: "Video Extraction cannot be done on more than 1 item",
+        });
+        return;
+      }
+      const element = nodes.values().next().value;
+      if (element instanceof VideoNode !== true) {
+        uiAccess.dialogError.show({
+          title: "Error: Video Extraction",
+          message:
+            "Extraction is only available for Videos, it is not avaliable for other Assets yet",
+        });
+        this.selectionManager.clearSelection();
+      }
+      // Gating done
+
+      //cast medianode to videonode
+      this.videoExtractionHandler.startVideoExtraction(element as VideoNode);
+    });
     uiEvents.toolbarNode.DELETE.onClick(() =>
       this.commandManager.deleteNodes(),
     );
@@ -386,23 +439,39 @@ export class Engine {
       this.addText(textdata);
     });
 
-   
-    let renderTimeout: NodeJS.Timeout;
-
-    uiEvents.promptEvents.onPromptStrengthChanged(async (strength) => {
-      this.realTimeDrawEngine.currentStrength = strength / 100.0;
-      clearTimeout(renderTimeout);
-      renderTimeout = setTimeout(async () => {
-        await this.realTimeDrawEngine.render();
-      }, 1000);
+    uiEvents.onChromakeyRequest((chromakeyProps) => {
+      const node = this.selectionManager
+        .getSelectedNodes()
+        .values()
+        .next().value;
+      if (!node) {
+        console.log("Node was not returned.");
+        return;
+      }
+      if (node instanceof VideoNode) {
+        if (chromakeyProps.isChromakeyEnabled) {
+          this.commandManager.addChromaKey({
+            videoNode: node,
+            newChromaColor: chromakeyProps.chromakeyColor ?? {
+              red: 120,
+              green: 150,
+              blue: 120,
+            },
+          });
+        } else {
+          this.commandManager.removeChromaKey({
+            videoNode: node,
+          });
+        }
+      }
     });
 
-    uiEvents.promptEvents.onPromptTextChanged(async (prompt) => {
+    uiEvents.promptEvents.onPromptStrengthChanged((strength) => {
+      this.realTimeDrawEngine.currentStrength = strength;
+    });
+
+    uiEvents.promptEvents.onPromptTextChanged((prompt) => {
       this.realTimeDrawEngine.currentPrompt = prompt;
-      clearTimeout(renderTimeout);
-      renderTimeout = setTimeout(async () => {
-        await this.realTimeDrawEngine.render();
-      }, 1000);
     });
 
     uiEvents.onAddShapeToEngine((shapeData) => {
@@ -419,11 +488,6 @@ export class Engine {
       }
     });
 
-
-    uiEvents.toolbarMain.onBgColorChanged((data)=>{
-      this.realTimeDrawEngine.updateBackground(data);
-    });
-
     uiEvents.toolbarNode.color.onConfirmChanged((nodeColor) => {
       console.log("Color change event triggered with:", nodeColor);
       if (!nodeColor) {
@@ -576,7 +640,6 @@ export class Engine {
   }
 
   public addShape(type: ShapeType, size: number, color?: string) {
-
     const shapeNode = new ShapeNode({
       canvasPosition: this.realTimeDrawEngine.captureCanvas.position(),
       canvasSize: this.realTimeDrawEngine.captureCanvas.size(),
@@ -587,9 +650,9 @@ export class Engine {
       selectionManagerRef: this.selectionManager,
       loaded: async () => {
         await this.realTimeDrawEngine.render();
-        this.setAppMode(AppModes.SELECT);
       },
     });
+    shapeNode.kNode.zIndex(1); // Replace desiredZIndex with a number
 
     this.commandManager.createNode(shapeNode);
 
@@ -604,7 +667,6 @@ export class Engine {
       selectionManagerRef: this.selectionManager,
       loaded: async () => {
         this.realTimeDrawEngine.render();
-        this.setAppMode(AppModes.SELECT);
       },
     });
     this.commandManager.createNode(copyNode);
@@ -618,10 +680,9 @@ export class Engine {
       selectionManagerRef: this.selectionManager,
       loaded: async () => {
         await this.realTimeDrawEngine.render();
-        this.setAppMode(AppModes.SELECT);
       },
     });
-
+    imageNode.kNode.zIndex(1);
     this.commandManager.createNode(imageNode);
   }
   public addPaintNode(

frontend/realtime/src/KonvaApp/Nodes/ImageNode.ts

@@ -106,7 +106,7 @@ export class ImageNode extends NetworkedNode {
       this.kNode.image(imageComponent);
       this.kNode.setSize(adjustedSize);
       this.kNode.setPosition(centerPosition);
-      this.kNode.moveToTop();
+
       this.kNode.fill(transparent);
       this.listenToBaseKNode();
       this.mediaLayerRef.draw();

frontend/realtime/src/KonvaApp/Nodes/PaintNode.ts

@@ -6,16 +6,16 @@ import { NodeType } from "./constants";
 import { NetworkedNode } from "./NetworkedNode";
 
 interface PaintNodeConstructor {
-  canvasElement: HTMLCanvasElement;
+  canvasElement:HTMLCanvasElement;
   lineBounds: {
     width: number;
     height: number;
     x: number;
     y: number;
-  };
+    };
   mediaLayerRef: Konva.Layer;
   selectionManagerRef: SelectionManager;
-
+  
   loaded: () => Promise<void>;
 }
 
@@ -31,7 +31,8 @@ export class PaintNode extends NetworkedNode {
     selectionManagerRef,
     loaded,
   }: PaintNodeConstructor) {
-    const cl = canvasElement;
+
+    const cl = canvasElement
 
     // Convert circle to image
     const dataURL = cl.toDataURL();
@@ -39,31 +40,31 @@ export class PaintNode extends NetworkedNode {
     imageObj.src = dataURL;
 
     const paintShape = new Konva.Image({
-      x: lineBounds.x,
-      y: lineBounds.y,
-      width: lineBounds.width,
-      height: lineBounds.height,
-      image: imageObj,
-      listening: false, // Disable event listening to allow captureCanvas to receive events
-      draggable: true,
-      // stroke: '#000000',
-      // strokeWidth: 2,
-      // dash: [5, 5], // Creates dotted outline THIS IS FOR DEBUG
-      globalCompositeOperation: "source-over",
-      fill: "transparent",
-    });
+        x: lineBounds.x,
+        y: lineBounds.y,
+        width: lineBounds.width,
+        height: lineBounds.height,
+        image: imageObj,
+        listening: false, // Disable event listening to allow captureCanvas to receive events
+        draggable: true,
+        // stroke: '#000000',
+        // strokeWidth: 2,
+        // dash: [5, 5], // Creates dotted outline THIS IS FOR DEBUG
+        globalCompositeOperation: 'source-over',
+        fill: 'transparent' 
+      });
 
     imageObj.onload = () => {
-      this.mediaLayerRef.draw();
-      loaded();
+        this.mediaLayerRef.draw();
+        loaded();
     };
-
+  
     super({
       selectionManagerRef: selectionManagerRef,
       mediaLayerRef: mediaLayerRef,
       kNode: paintShape,
     });
-    this.kNode.moveToTop();
+
     this.mediaLayerRef.add(this.kNode);
     this.listenToBaseKNode();
     this.mediaLayerRef.draw();

frontend/realtime/src/KonvaApp/Nodes/PreviewCopy.ts

@@ -30,7 +30,7 @@ export class PreviewCopyNode extends NetworkedNode {
       mediaLayerRef: mediaLayerRef,
       kNode: image,
     });
-    this.kNode.moveToTop();
+
     this.mediaLayerRef.add(this.kNode);
     this.listenToBaseKNode();
     this.mediaLayerRef.draw();

frontend/realtime/src/KonvaApp/Nodes/ShapeNode.ts

@@ -117,7 +117,7 @@ export class ShapeNode extends NetworkedNode {
 
         squareImageObj.onload = () => {
           this.mediaLayerRef.draw();
-          loaded();
+          loaded()
         };
 
         // Clean up the temporary square
@@ -153,7 +153,7 @@ export class ShapeNode extends NetworkedNode {
 
         triangleImageObj.onload = () => {
           this.mediaLayerRef.draw();
-          loaded();
+          loaded()
         };
 
         // Clean up the temporary triangle
@@ -180,8 +180,6 @@ export class ShapeNode extends NetworkedNode {
       componentSize: canvasSize,
       maxSize: canvasSize,
     });
-    // bring to the front
-    this.kNode.moveToTop();
     this.kNode.setPosition(centerPosition);
     this.mediaLayerRef.add(this.kNode);
     this.listenToBaseKNode();

frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts

@@ -15,24 +15,14 @@ import { Image } from "@tauri-apps/api/image";
 
 import { PaintNode } from "../Nodes/PaintNode";
 
-import { WebSocketClient } from "../../PyServer/DrawServer";
-import { PreviewCopyNode } from "../Nodes/PreviewCopy";
+import {
+  ServerSetupPayload,
+  ServerSettingsPayload,
+  ServerResponse,
+} from "../types/ServerTypes";
 
 // https://www.aiseesoft.com/resource/phone-aspect-ratio-screen-resolution.html#:~:text=16%3A9%20Aspect%20Ratio
 
-interface ServerSettings {
-  model_path: string;
-  lora_path?: string;
-}
-
-interface GenerateImageParams {
-  image: string;
-  prompt: string;
-  strength?: number;
-  guidance_scale?: number;
-  num_inference_steps?: number;
-}
-
 export class RealTimeDrawEngine {
   private videoNodes: VideoNode[];
   private imageNodes: (ImageNode | TextNode | ShapeNode | PaintNode)[];
@@ -70,7 +60,7 @@ export class RealTimeDrawEngine {
   // paint Brush Size
   // has to exit out of paint mode when shape or image are used.
   public paintColor: string = "#000000";
-  private brushSize: number = 5; // Default brush size
+  private brushSize: number = 5;  // Default brush size
 
   private onDrawCallback?: (
     canvas: HTMLCanvasElement,
@@ -84,12 +74,8 @@ export class RealTimeDrawEngine {
 
   private onPreviewCopyCallback?: (previewCopy: Konva.Image) => void; // New Callback
 
-  private client: WebSocketClient | null = null;
-  private isConnected: boolean = false;
+  private serverSocket: WebSocket | null = null;
 
-
-  private backgroundNode: ShapeNode | null = null;
-  public backgroundColor: string = "#d2d2d2"
   constructor({
     width,
     height,
@@ -141,7 +127,6 @@ export class RealTimeDrawEngine {
     this.drawingsLayer = new Konva.Layer({
       clearBeforeDraw: true, // Ensures transparent background
     });
- 
     this.mediaLayerRef.getStage()?.add(this.drawingsLayer); // to od pass in stage
 
     // Set background layer to red and media layer to green for visibility
@@ -151,7 +136,7 @@ export class RealTimeDrawEngine {
     this.fps = 24;
 
     this.currentPrompt = "";
-    this.currentStrength = 1;
+    this.currentStrength = 100;
     // This is captures a subset of the medialayer ref
     this.captureCanvas = new Konva.Rect({
       name: "CaptureCanvas",
@@ -159,7 +144,7 @@ export class RealTimeDrawEngine {
       y: this.positionY,
       width: this.width,
       height: this.height,
-      fill: this.backgroundColor,
+      fill: "white",
       stroke: "black",
       strokeWidth: 1,
       draggable: false,
@@ -186,104 +171,102 @@ export class RealTimeDrawEngine {
     // Add mouse events for preview canvas copying
     this.previewCopyListener();
 
-    this.startServer();
+    //this.startServer();
   }
 
   private isEnabled: boolean = false;
   private cleanupFunction: (() => void) | null = null;
 
   // this starts the python server
-  public async startServer() {
-    try {
-      this.client = new WebSocketClient("ws://localhost:8765");
+  public startServer() {
+    // Create WebSocket connection
+    const socket = new WebSocket("ws://localhost:8765");
 
-      // Set up message handlers
-      this.client.onProgress((progress) => {
-        console.log(
-          `Progress: ${progress.message} (${progress.progress * 100}%)`,
-        );
-        if (progress.error) {
-          console.error("Error:", progress.error);
+    // Setup event handlers
+    socket.onopen = () => {
+      console.log("Connected to inference server");
+      // Send initial setup payload to load models
+      const setupPayload: ServerSetupPayload = {
+        type: "setup",
+        model: {
+          name: "stabilityai/sdxl-turbo",
+          path: "F:/ComfyUI_windows_portable_nvidia/ComfyUI_windows_portable/ComfyUI/models/checkpoints/ponyDiffusionV6XL_v6StartWithThisOne.safetensors",
+          precision: "fp16",
+        },
+        lora: {
+          path: "F:/ComfyUI_windows_portable_nvidia/ComfyUI_windows_portable/ComfyUI/models/loras/LCM_LoRA_Weights_SDXL.safetensors",
+          alpha: 0.75,
+        },
+        device: "cuda",
+      };
+
+      socket.send(JSON.stringify(setupPayload));
+    };
+
+    // Handle incoming messages from server
+    socket.onmessage = (event) => {
+      try {
+        const response: ServerResponse = JSON.parse(event.data);
+        console.log("Server response:", response);
+        // Handle model loading progress updates
+        if (response.type === "progress") {
+          console.log(`Loading progress: ${response.percent}%`);
+          // You might want to update UI with loading progress here
         }
-      });
-
-      this.client.onResult(async (response) => {
-        try {
-          // Convert the base64 image to ImageBitmap and update the preview
-          const bitmap = await this.base64ToImageBitmap(response.image);
-          this.outputBitmap = bitmap;
-          this.previewCanvas.image(bitmap);
-          this.mediaLayerRef.batchDraw();
-        } catch (error) {
-          console.error("Error processing result image:", error);
-        } finally {
-          this.isProcessing = false;
+        // Handle successful model load
+        else if (
+          response.type === "success" &&
+          response.status === "model_loaded"
+        ) {
+          console.log("Model loaded successfully");
+          // Enable UI elements that depend on model being loaded
         }
-      });
+        // Handle image generation results
+        else if (
+          response.type === "success" &&
+          response.status === "generation_complete"
+        ) {
+          if (response.image) {
+            // Convert base64 image to ImageBitmap for preview
+            this.base64ToImageBitmap(response.image).then((bitmap) => {
+              this.outputBitmap = bitmap;
+              this.previewCanvas.image(bitmap);
+              this.mediaLayerRef.draw();
+            });
+          }
+        }
+      } catch (error) {
+        console.error("Error parsing server response:", error);
+      }
+    };
 
-      // Wait for connection to be established
-      await new Promise<void>((resolve, reject) => {
-        if (!this.client) return reject("Client not initialized");
+    socket.onerror = (error) => {
+      console.error("WebSocket error:", error);
+    };
 
-        this.client.ws.onopen = () => {
-          this.isConnected = true;
-          resolve();
-        };
-        this.client.ws.onerror = (error) => reject(error);
-      });
+    socket.onclose = () => {
+      console.log("Disconnected from inference server");
+    };
 
-      // Load initial model
-      await this.loadModel({
-        model_path: "C:/Users/Tensor/Downloads/animagineXL40_v4Opt.safetensors",
-        lora_path: "C:/Users/Tensor/Downloads/anyu_all.safetensors",
-      });
-    } catch (error) {
-      console.error("Failed to start server:", error);
-      this.isConnected = false;
-    }
+    // Store socket reference for later use
+    this.serverSocket = socket;
   }
 
-  private async loadModel(settings: ServerSettings) {
-    if (!this.client || !this.isConnected) {
-      throw new Error("Server not connected");
+  // Add method to update model settings
+  public updateServerSettings(settings: ServerSettingsPayload) {
+    if (this.serverSocket && this.serverSocket.readyState === WebSocket.OPEN) {
+      this.serverSocket.send(JSON.stringify(settings));
+    } else {
+      console.error(
+        "Cannot update settings: Server connection not established",
+      );
     }
-    await this.client.loadModel(settings);
-  }
-
-  private updateCursor(stage: Konva.Stage) {
-    // Create cursor canvas
-    const cursorCanvas = document.createElement('canvas');
-    const size = this.brushSize * 5; // Match the brush size used in drawing
-    cursorCanvas.width = size * 2;  // Double size for padding
-    cursorCanvas.height = size * 2;
-    
-    const ctx = cursorCanvas.getContext('2d');
-    if (!ctx) return;
-
-    // Draw the circle
-    ctx.beginPath();
-    ctx.arc(size, size, size/2, 0, Math.PI * 2);
-    ctx.strokeStyle = 'black';
-    ctx.lineWidth = 1;
-    ctx.stroke();
-
-    // Convert to data URL
-    const cursorUrl = cursorCanvas.toDataURL();
-    
-    // Apply custom cursor
-    stage.container().style.cursor = `url(${cursorUrl}) ${size} ${size}, auto`;
   }
 
   public paintMode() {
     let isDrawing = false;
     let currentLine: Konva.Line | null = null;
 
-    const stage = this.mediaLayerRef.getStage();
-    if (!stage) return;
-
-    // Initialize cursor
-    this.updateCursor(stage);
-
     const startDrawing = (pos: { x: number; y: number }) => {
       if (!this.isEnabled) return;
       const stage = this.mediaLayerRef.getStage();
@@ -304,7 +287,7 @@ export class RealTimeDrawEngine {
         y: this.captureCanvas.y(),
         draggable: false,
       });
-      this.drawingsLayer.moveToTop();
+
       this.drawingsLayer.add(currentLine); // Add to drawingsLayer
       isDrawing = true;
     };
@@ -377,6 +360,9 @@ export class RealTimeDrawEngine {
     };
 
     // Add event listeners
+    const stage = this.mediaLayerRef.getStage();
+    if (!stage) return;
+
     stage.on("mousedown touchstart", (e) => {
       const pos = stage.getPointerPosition();
       if (pos && isWithinCaptureCanvas(pos)) {
@@ -406,26 +392,17 @@ export class RealTimeDrawEngine {
 
   public enablePaintMode() {
     this.isEnabled = true;
-    this.disableDragging();
     if (!this.cleanupFunction) {
       this.paintMode();
     }
-    // Update cursor when enabling paint mode
-    const stage = this.mediaLayerRef.getStage();
-    if (stage) {
-      this.updateCursor(stage);
-    }
   }
 
   public enableDragging() {
     // Enable dragging for all nodes in media layer
     this.imageNodes?.forEach((node) => {
-      console.log("Image Nodes Enable");
-      console.log(node);
       node.kNode.draggable(true);
       node.kNode.listening(true);
     });
-
     this.mediaLayerRef.batchDraw();
   }
 
@@ -440,96 +417,89 @@ export class RealTimeDrawEngine {
 
   public disablePaintMode() {
     this.isEnabled = false;
-    this.enableDragging();
     if (this.cleanupFunction) {
       this.cleanupFunction();
       this.cleanupFunction = null;
     }
-    // Reset cursor when disabling paint mode
-    const stage = this.mediaLayerRef.getStage();
-    if (stage) {
-      stage.container().style.cursor = 'default';
-    }
   }
 
   public previewCopyListener() {
-    // Start of Selection
-    this.previewCanvas.on("mousedown touchstart", () => {
-      if (!this.outputBitmap) {
-        console.log("No preview image to copy");
-        return;
-      }
-
-      // Create draggable preview copy
-      const previewCopy = new Konva.Image({
-        x: this.previewCanvas.x(),
-        y: this.previewCanvas.y(),
-        width: this.width,
-        height: this.height,
-        image: this.outputBitmap,
-        draggable: true,
-        listening: true,
-      });
-
-      if (this.onPreviewCopyCallback) {
-        this.onPreviewCopyCallback(previewCopy);
-      }
-      previewCopy.startDrag();
-      previewCopy.moveToTop();
-      previewCopy.on("dragend", (e: Konva.KonvaEventObject<DragEvent>) => {
-        const previewBox = previewCopy.getClientRect();
-        const captureBox = this.captureCanvas.getClientRect();
-
-        if (Konva.Util.haveIntersection(previewBox, captureBox)) {
-          // Snap to capture canvas position
-          previewCopy.position({
-            x: this.captureCanvas.x(),
-            y: this.captureCanvas.y(),
+        // Start of Selection
+        this.previewCanvas.on("mousedown touchstart", () => {
+          if (!this.outputBitmap) {
+            console.log("No preview image to copy");
+            return;
+          }
+    
+          // Create draggable preview copy
+          const previewCopy = new Konva.Image({
+            x: this.previewCanvas.x(),
+            y: this.previewCanvas.y(),
+            width: this.width,
+            height: this.height,
+            image: this.outputBitmap,
+            draggable: true,
+            listening: true,
           });
-
-          e.target.off("dragend");
-        } else {
+    
+          if (this.onPreviewCopyCallback) {
+            this.onPreviewCopyCallback(previewCopy);
+          }
+          previewCopy.startDrag();
+          previewCopy.on("dragend", (e: Konva.KonvaEventObject<DragEvent>) => {
+            const previewBox = previewCopy.getClientRect();
+            const captureBox = this.captureCanvas.getClientRect();
+    
+            if (Konva.Util.haveIntersection(previewBox, captureBox)) {
+              // Snap to capture canvas position
+              previewCopy.position({
+                x: this.captureCanvas.x(),
+                y: this.captureCanvas.y(),
+              });
+              e.target.off('dragend');
+              previewCopy.setZIndex(1); // send back after drop.
+            } else {
+                // Start Generation Here
+                previewCopy.destroy();
+                this.mediaLayerRef.batchDraw();
+            }
+          });
+          
           // Start Generation Here
-          previewCopy.destroy();
-          this.mediaLayerRef.batchDraw();
-        }
-      });
+          const stage = this.mediaLayerRef.getStage();
+          const layer = this.mediaLayerRef; // Use the existing media layer
 
-      // Start Generation Here
-      const stage = this.mediaLayerRef.getStage();
-      const layer = this.mediaLayerRef; // Use the existing media layer
+          // Create centered text
+              // Start of Selection
+              const centeredText = new Konva.Text({
+                text: 'Hold and Drag Over, To Copy',
+                fontSize: 24,
+                fontFamily: 'Arial',
+                fill: 'black',
+                x: this.captureCanvas.x() + this.captureCanvas.width() / 2,
+                y: this.captureCanvas.y() + this.captureCanvas.height() / 2,
+                id: 'copyText',
+                listening: false,
+              });
+              // Calculate and set offsets based on text size to center it
+              centeredText.offsetX(centeredText.width() / 2);
+              centeredText.offsetY(centeredText.height() / 2);
 
-      // Create centered text
-      // Start of Selection
-      const centeredText = new Konva.Text({
-        text: "Hold and Drag Over, To Copy",
-        fontSize: 24,
-        fontFamily: "Arial",
-        fill: "black",
-        x: this.captureCanvas.x() + this.captureCanvas.width() / 2,
-        y: this.captureCanvas.y() + this.captureCanvas.height() / 2,
-        id: "copyText",
-        listening: false,
-      });
-      // Calculate and set offsets based on text size to center it
-      centeredText.offsetX(centeredText.width() / 2);
-      centeredText.offsetY(centeredText.height() / 2);
-
-      layer.add(centeredText);
-      centeredText.moveToTop();
-      layer.batchDraw();
-
-      // Remove text on mouseup/touchend
-      const removeText = () => {
-        const text = layer.findOne("#copyText");
-        if (text) {
-          text.destroy();
+          layer.add(centeredText);
+          centeredText.moveToTop();
           layer.batchDraw();
-        }
-        stage.off("mouseup touchend", removeText);
-      };
-      stage.on("mouseup touchend", removeText);
-    });
+    
+          // Remove text on mouseup/touchend
+          const removeText = () => {
+            const text = layer.findOne('#copyText');
+            if (text) {
+              text.destroy();
+              layer.batchDraw();
+            }
+            stage.off("mouseup touchend", removeText);
+          };
+          stage.on("mouseup touchend", removeText);
+        });
   }
 
   public findImageNodeById(
@@ -677,7 +647,12 @@ export class RealTimeDrawEngine {
   public isProcessing = false;
   private handleNodeDragEnd = async () => {
     // Clean up any existing state
-
+    if (this.isProcessing) {
+      console.log("isProcessing Returning");
+      return;
+    }
+    console.log("Node drag ended");
+    this.isProcessing = true;
     await this.render();
   };
 
@@ -686,20 +661,19 @@ export class RealTimeDrawEngine {
       node instanceof ImageNode ||
       node instanceof TextNode ||
       node instanceof ShapeNode ||
-      node instanceof PaintNode ||
-      node instanceof PreviewCopyNode
+      node instanceof PaintNode
     ) {
       console.debug("Adding node:", node);
       this.imageNodes.push(node);
       console.log(this.imageNodes);
+      //node.kNode.on("dragend", this.handleNodeDragEnd);
     }
 
     // ensure the layer doesn't move if added while painting.
     if (this.isEnabled) {
       this.disableDragging();
     }
-    // ensure this goes to the top.
-    this.previewCanvas.moveToTop();
+
     await this.render();
   }
 
@@ -714,8 +688,7 @@ export class RealTimeDrawEngine {
       node instanceof ImageNode ||
       node instanceof TextNode ||
       node instanceof ShapeNode ||
-      node instanceof PaintNode ||
-      node instanceof PreviewCopyNode
+      node instanceof PaintNode
     ) {
       const index = this.imageNodes.indexOf(node);
       if (index > -1) {
@@ -876,12 +849,6 @@ export class RealTimeDrawEngine {
 
   public async render() {
     // only pick nodes that intersect wi th the canvas on screen bounds to freeze.
-    if (this.isProcessing) {
-      console.log("isProcessing Returning");
-      return;
-    }
-    console.log("Calling Render");
-    this.isProcessing = true;
 
     this.mediaLayerRef.draw();
     // Output all nodes in mediaLayerRef
@@ -905,20 +872,9 @@ export class RealTimeDrawEngine {
 
     // Test code
     if (true) {
-      //this.outputBitmap = bitmap;
-      //this.previewCanvas.image(bitmap);
-      const base64Bitmap = await this.imageBitmapToBase64(bitmap);
-      console.log(this.currentStrength);
-      await this.client?.generateImage({
-        image: base64Bitmap,
-        prompt: this.currentPrompt,
-        strength: this.currentStrength,
-        guidance_scale: 2,
-        num_inference_steps: 4,
-        height: 1024,
-        width: 1024,
-        lora_strength: 1.0,
-      });
+      this.outputBitmap = bitmap;
+      this.previewCanvas.image(bitmap);
+      this.isProcessing = false;
       return;
     }
 
@@ -947,33 +903,10 @@ export class RealTimeDrawEngine {
 
   // Add getter/setter for brush size
   public set paintBrushSize(size: number) {
-    this.brushSize =  size * 7.5;
-    const stage = this.mediaLayerRef.getStage();
-    if (stage && this.isEnabled) {
-      this.updateCursor(stage);
-    }
+    this.brushSize = size;
   }
 
   public get paintBrushSize(): number {
     return this.brushSize;
   }
-
-  // Add method to check connection status
-  public isServerConnected(): boolean {
-    return this.isConnected && this.client !== null;
-  }
-
-  // Add method to reconnect if needed
-  public async reconnect() {
-    if (this.client) {
-      this.client.ws.close();
-    }
-    await this.startServer();
-  }
-
-  // Add method to create or update background
-  public updateBackground(color: string) {
-      this.captureCanvas.fill(color);
-      this.mediaLayerRef.batchDraw();
-  }
 }

frontend/realtime/src/PyServer/DrawServer.ts

@@ -1,67 +0,0 @@
-// websocket-client.ts
-import {
-  ProgressMessage,
-  GenerateResponse,
-  LoadModelRequest,
-  GenerateRequest,
-} from "./DrawServerTypes";
-
-export class WebSocketClient {
-  public ws: WebSocket;
-  private messageHandlers: Map<string, (data: any) => void>;
-  private isConnected: boolean = false;
-
-  constructor(url: string = "ws://localhost:8765") {
-    this.ws = new WebSocket(url);
-    this.messageHandlers = new Map();
-
-    this.ws.onmessage = (event) => this.handleMessage(event);
-    this.ws.onerror = (error) => console.error("WebSocket error:", error);
-    this.ws.onopen = () => this.isConnected = true;
-    this.ws.onclose = () => this.isConnected = false;
-  }
-
-  private handleMessage(event: MessageEvent) {
-    const data = JSON.parse(event.data);
-    const handler = this.messageHandlers.get(data.type);
-    console.log(data);
-    if (handler) {
-      handler(data);
-    }
-  }
-
-  public onProgress(callback: (message: ProgressMessage) => void) {
-    this.messageHandlers.set("progress", callback);
-  }
-
-  public onResult(callback: (response: GenerateResponse) => void) {
-    this.messageHandlers.set("result", callback);
-  }
-
-  public async loadModel(
-    request: Omit<LoadModelRequest, "command">,
-  ): Promise<void> {
-    const message: LoadModelRequest = {
-      command: "load_model",
-      ...request,
-    };
-    this.ws.send(JSON.stringify(message));
-  }
-
-  public async generateImage(
-    request: Omit<GenerateRequest, "command">,
-  ): Promise<void> {
-    const message: GenerateRequest = {
-      command: "generate",
-      strength: 0.6,
-      guidance_scale: 2.0,
-      num_inference_steps: 4,
-      ...request,
-    };
-    this.ws.send(JSON.stringify(message));
-  }
-
-  public isServerConnected(): boolean {
-    return this.isConnected && this.ws.readyState === WebSocket.OPEN;
-  }
-}

frontend/realtime/src/PyServer/DrawServerTypes.ts

@@ -1,83 +0,0 @@
-// types.ts
-export interface ProgressMessage {
-  type: "progress";
-  message: string;
-  progress: number;
-  error?: string | null;
-}
-
-export interface LoadModelRequest {
-  command: "load_model";
-  model_path: string;
-  lora_path?: string;
-}
-
-export interface GenerateRequest {
-  command: "generate";
-  image: string;
-  prompt: string;
-  strength?: number;
-  guidance_scale?: number;
-  num_inference_steps?: number;
-  height:number;
-  width:number;
-  lora_strength:number;
-}
-
-export interface GenerateResponse {
-  type: "result";
-  image: string;
-}
-
-// websocket-client.ts
-export class WebSocketClient {
-  private ws: WebSocket;
-  private messageHandlers: Map<string, (data: any) => void>;
-
-  constructor(url: string = "ws://localhost:8765") {
-    this.ws = new WebSocket(url);
-    this.messageHandlers = new Map();
-
-    this.ws.onmessage = (event) => this.handleMessage(event);
-    this.ws.onerror = (error) => console.error("WebSocket error:", error);
-  }
-
-  private handleMessage(event: MessageEvent) {
-    const data = JSON.parse(event.data);
-    const handler = this.messageHandlers.get(data.type);
-    if (handler) {
-      handler(data);
-    }
-  }
-
-  public onProgress(callback: (message: ProgressMessage) => void) {
-    this.messageHandlers.set("progress", callback);
-  }
-
-  public onResult(callback: (response: GenerateResponse) => void) {
-    this.messageHandlers.set("result", callback);
-  }
-
-  public async loadModel(
-    request: Omit<LoadModelRequest, "command">,
-  ): Promise<void> {
-    const message: LoadModelRequest = {
-      command: "load_model",
-      ...request,
-    };
-    this.ws.send(JSON.stringify(message));
-  }
-
-  public async generateImage(
-    request: Omit<GenerateRequest, "command">,
-  ): Promise<void> {
-    const message: GenerateRequest = {
-      command: "generate",
-      strength: 0.6,
-      guidance_scale: 2.0,
-      num_inference_steps: 4,
-      ...request,
-    };
-    this.ws.send(JSON.stringify(message));
-  }
-}

frontend/realtime/src/PyServer/ServerClient.ts

@@ -0,0 +1,260 @@
+import {
+  ServerMessage,
+  WebSocketHandlers,
+  SetupRequest,
+  UpdateSettingsRequest,
+  GenerateRequest,
+  StatusRequest,
+  GenerationParameters,
+  ServerSettings,
+  createSetupRequest,
+  createUpdateSettingsRequest,
+  createGenerateRequest,
+  createStatusRequest,
+} from "./ServerTypes";
+
+/**
+ * ServerClient provides a typed interface for communicating with the Python WebSocket server
+ */
+export class ServerClient {
+  private ws: WebSocket | null = null;
+  private handlers: WebSocketHandlers = {};
+  private connectionPromise: Promise<void> | null = null;
+  private connectionResolve: (() => void) | null = null;
+  private connectionReject: ((error: Error) => void) | null = null;
+
+  /**
+   * Create a new ServerClient
+   * @param url The WebSocket URL to connect to
+   * @param handlers Optional event handlers for different message types
+   */
+  constructor(
+    private url: string,
+    handlers: Partial<WebSocketHandlers> = {},
+  ) {
+    this.handlers = { ...handlers };
+  }
+
+  /**
+   * Connect to the WebSocket server
+   * @returns A promise that resolves when connected
+   */
+  public connect(): Promise<void> {
+    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
+      return Promise.resolve();
+    }
+
+    if (this.connectionPromise) {
+      return this.connectionPromise;
+    }
+
+    this.connectionPromise = new Promise<void>((resolve, reject) => {
+      this.connectionResolve = resolve;
+      this.connectionReject = reject;
+
+      try {
+        this.ws = new WebSocket(this.url);
+
+        this.ws.onopen = () => {
+          if (this.connectionResolve) {
+            this.connectionResolve();
+          }
+          if (this.handlers.onOpen) {
+            this.handlers.onOpen();
+          }
+        };
+
+        this.ws.onclose = () => {
+          if (this.handlers.onClose) {
+            this.handlers.onClose();
+          }
+          this.ws = null;
+          this.connectionPromise = null;
+        };
+
+        this.ws.onerror = (event) => {
+          const error = new Error(`WebSocket error: ${event}`);
+          if (this.connectionReject) {
+            this.connectionReject(error);
+          }
+          if (this.handlers.onError) {
+            this.handlers.onError({
+              type: "error",
+              error: `Connection error: ${event}`,
+            });
+          }
+        };
+
+        this.ws.onmessage = (event) => {
+          try {
+            const message = JSON.parse(event.data) as ServerMessage;
+            this.handleMessage(message);
+          } catch (error) {
+            console.error("Error parsing message:", error);
+            if (this.handlers.onError) {
+              this.handlers.onError({
+                type: "error",
+                error: `Failed to parse message: ${error}`,
+              });
+            }
+          }
+        };
+      } catch (error) {
+        if (this.connectionReject) {
+          this.connectionReject(error as Error);
+        }
+        this.connectionPromise = null;
+      }
+    });
+
+    return this.connectionPromise;
+  }
+
+  /**
+   * Disconnect from the WebSocket server
+   */
+  public disconnect(): void {
+    if (this.ws) {
+      this.ws.close();
+      this.ws = null;
+    }
+    this.connectionPromise = null;
+  }
+
+  /**
+   * Send a message to the server
+   * @param message The message to send
+   * @returns A promise that resolves when the message is sent
+   */
+  private async send(message: ServerMessage): Promise<void> {
+    await this.connect();
+
+    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
+      throw new Error("WebSocket is not connected");
+    }
+
+    this.ws.send(JSON.stringify(message));
+  }
+
+  /**
+   * Handle an incoming message from the server
+   * @param message The message received from the server
+   */
+  private handleMessage(message: ServerMessage): void {
+    switch (message.type) {
+      case "setup_response":
+        if (this.handlers.onSetupResponse) {
+          this.handlers.onSetupResponse(message);
+        }
+        break;
+
+      case "update_settings_response":
+        if (this.handlers.onUpdateSettingsResponse) {
+          this.handlers.onUpdateSettingsResponse(message);
+        }
+        break;
+
+      case "generate_response":
+        if (this.handlers.onGenerateResponse) {
+          this.handlers.onGenerateResponse(message);
+        }
+        break;
+
+      case "status_response":
+        if (this.handlers.onStatusResponse) {
+          this.handlers.onStatusResponse(message);
+        }
+        break;
+
+      case "progress_update":
+        if (this.handlers.onProgressUpdate) {
+          this.handlers.onProgressUpdate(message);
+        }
+        break;
+
+      case "error":
+      case "connection_error":
+        if (this.handlers.onError) {
+          this.handlers.onError(message);
+        }
+        break;
+    }
+  }
+
+  /**
+   * Set event handlers
+   * @param handlers The event handlers to set
+   */
+  public setHandlers(handlers: Partial<WebSocketHandlers>): void {
+    this.handlers = { ...this.handlers, ...handlers };
+  }
+
+  /**
+   * Send a setup request to initialize the model
+   * @param sdxl_checkpoint_path Path to the SDXL checkpoint
+   * @param lora_path Optional path to a LoRA file
+   */
+  public async setup(
+    sdxl_checkpoint_path: string,
+    lora_path?: string,
+  ): Promise<void> {
+    const request = createSetupRequest(sdxl_checkpoint_path, lora_path);
+    await this.send(request);
+  }
+
+  /**
+   * Send an update settings request
+   * @param settings The settings to update
+   */
+  public async updateSettings(
+    settings: Partial<ServerSettings>,
+  ): Promise<void> {
+    const request = createUpdateSettingsRequest(settings);
+    await this.send(request);
+  }
+
+  /**
+   * Send a generate request
+   * @param params The generation parameters
+   */
+  public async generate(params: GenerationParameters): Promise<void> {
+    const request = createGenerateRequest(params);
+    await this.send(request);
+  }
+
+  /**
+   * Send a status request
+   */
+  public async requestStatus(): Promise<void> {
+    const request = createStatusRequest();
+    await this.send(request);
+  }
+}
+
+/**
+ * Create a base64 string from an image
+ * @param file The image file
+ * @returns A promise that resolves to the base64 string
+ */
+export function fileToBase64(file: File): Promise<string> {
+  return new Promise((resolve, reject) => {
+    const reader = new FileReader();
+    reader.readAsDataURL(file);
+    reader.onload = () => {
+      const result = reader.result as string;
+      // Remove the data URL prefix (e.g., "data:image/png;base64,")
+      const base64 = result.split(",")[1];
+      resolve(base64);
+    };
+    reader.onerror = reject;
+  });
+}
+
+/**
+ * Create an image from a base64 string
+ * @param base64 The base64 string
+ * @returns The image URL
+ */
+export function base64ToImageUrl(base64: string): string {
+  return `data:image/png;base64,${base64}`;
+}

frontend/realtime/src/PyServer/ServerTypes.ts

@@ -0,0 +1,216 @@
+// Server Message Types - Common base type for all messages
+export type MessageType =
+  | "setup"
+  | "setup_response"
+  | "update_settings"
+  | "update_settings_response"
+  | "generate"
+  | "generate_response"
+  | "status"
+  | "status_response"
+  | "progress_update"
+  | "error"
+  | "connection_error";
+
+// Base interface for all messages
+interface BaseMessage {
+  type: MessageType;
+}
+
+// ===== Setup Messages =====
+
+// Setup Request
+export interface SetupRequest extends BaseMessage {
+  type: "setup";
+  setup: {
+    sdxl_checkpoint_path: string;
+    lora_path?: string;
+  };
+}
+
+// Setup Response
+export interface SetupResponse extends BaseMessage {
+  type: "setup_response";
+  success: boolean;
+  message?: string;
+  error?: string;
+  settings?: ServerSettings;
+}
+
+// ===== Update Settings Messages =====
+
+// Default generation settings that can be configured
+export interface GenerationSettings {
+  default_lora_strength?: number;
+  default_image_to_image_strength?: number;
+  default_width?: number;
+  default_height?: number;
+  default_steps?: number;
+  default_guidance_scale?: number;
+}
+
+// Server settings - includes both paths and generation defaults
+export interface ServerSettings extends GenerationSettings {
+  sdxl_checkpoint_path: string;
+  lora_path?: string | null;
+}
+
+// Update Settings Request
+export interface UpdateSettingsRequest extends BaseMessage {
+  type: "update_settings";
+  settings: Partial<ServerSettings>;
+}
+
+// Update Settings Response
+export interface UpdateSettingsResponse extends BaseMessage {
+  type: "update_settings_response";
+  success: boolean;
+  message?: string;
+  error?: string;
+  updates_applied?: string[];
+  current_settings: ServerSettings;
+}
+
+// ===== Generate Messages =====
+
+// Generation parameters
+export interface GenerationParameters {
+  image: string; // base64 encoded image
+  prompt: string;
+  lora_strength?: number;
+  image_to_image_strength?: number;
+  generated_image_width?: number;
+  generated_image_height?: number;
+  num_inference_steps?: number;
+  guidance_scale?: number;
+  seed?: number;
+}
+
+// Generate Request
+export interface GenerateRequest extends BaseMessage {
+  type: "generate";
+  generate: GenerationParameters;
+}
+
+// Generate Response
+export interface GenerateResponse extends BaseMessage {
+  type: "generate_response";
+  success: boolean;
+  image?: string; // base64 encoded image
+  error?: string;
+}
+
+// ===== Status Messages =====
+
+// Status Request
+export interface StatusRequest extends BaseMessage {
+  type: "status";
+}
+
+// Status Response
+export interface StatusResponse extends BaseMessage {
+  type: "status_response";
+  success: boolean;
+  model_initialized: boolean;
+  current_settings: ServerSettings;
+}
+
+// ===== Progress Update =====
+
+// Progress stages
+export type ProgressStage =
+  | "downloading"
+  | "loading_unet"
+  | "loading_base_model"
+  | "configuring_scheduler"
+  | "loading_lcm_adapter"
+  | "loading_lora"
+  | "generating"
+  | "error";
+
+// Progress Update Message
+export interface ProgressUpdateMessage extends BaseMessage {
+  type: "progress_update";
+  stage: ProgressStage;
+  progress: number; // 0-100
+  file: string;
+}
+
+// ===== Error Messages =====
+
+// Error Message
+export interface ErrorMessage extends BaseMessage {
+  type: "error";
+  error: string;
+}
+
+// Connection Error Message
+export interface ConnectionErrorMessage extends BaseMessage {
+  type: "connection_error";
+  error: string;
+}
+
+// Union type for all possible server messages
+export type ServerMessage =
+  | SetupRequest
+  | SetupResponse
+  | UpdateSettingsRequest
+  | UpdateSettingsResponse
+  | GenerateRequest
+  | GenerateResponse
+  | StatusRequest
+  | StatusResponse
+  | ProgressUpdateMessage
+  | ErrorMessage
+  | ConnectionErrorMessage;
+
+// WebSocket event handlers for different message types
+export interface WebSocketHandlers {
+  onSetupResponse?: (response: SetupResponse) => void;
+  onUpdateSettingsResponse?: (response: UpdateSettingsResponse) => void;
+  onGenerateResponse?: (response: GenerateResponse) => void;
+  onStatusResponse?: (response: StatusResponse) => void;
+  onProgressUpdate?: (update: ProgressUpdateMessage) => void;
+  onError?: (error: ErrorMessage | ConnectionErrorMessage) => void;
+  onOpen?: () => void;
+  onClose?: () => void;
+}
+
+// Helper function to create a typed setup request
+export function createSetupRequest(
+  sdxl_checkpoint_path: string,
+  lora_path?: string,
+): SetupRequest {
+  return {
+    type: "setup",
+    setup: {
+      sdxl_checkpoint_path,
+      lora_path,
+    },
+  };
+}
+
+// Helper function to create a typed update settings request
+export function createUpdateSettingsRequest(
+  settings: Partial<ServerSettings>,
+): UpdateSettingsRequest {
+  return {
+    type: "update_settings",
+    settings,
+  };
+}
+
+// Helper function to create a typed generate request
+export function createGenerateRequest(
+  params: GenerationParameters,
+): GenerateRequest {
+  return {
+    type: "generate",
+    generate: params,
+  };
+}
+
+// Helper function to create a typed status request
+export function createStatusRequest(): StatusRequest {
+  return { type: "status" };
+}

frontend/realtime/src/app/index.tsx

@@ -5,7 +5,6 @@ import { RouterProvider } from "react-router-dom";
 import { router } from "./router";
 import { posthog } from "posthog-js";
 import { PostHogProvider } from "posthog-js/react";
-import { Toaster } from "~/components/ui/Toast";
 
 import "./global.css";
 import { useRenderCounter } from "~/hooks/useRenderCounter";
@@ -32,17 +31,11 @@ const App = () => {
       <StrictMode>
         <PostHogProvider client={posthog}>
           <RouterProvider router={router} />
-          <Toaster />
         </PostHogProvider>
       </StrictMode>
     );
   }
-  return (
-    <>
-      <RouterProvider router={router} />
-      <Toaster />
-    </>
-  );
+  return <RouterProvider router={router} />;
 };
 
 createRoot(document.getElementById("root")!).render(<App />);

frontend/realtime/src/components/ui/Toast.tsx

@@ -1,12 +0,0 @@
-import { Toaster as SonnerToaster } from "sonner";
-
-export const Toaster = () => {
-  return (
-    <SonnerToaster
-      position="top-right"
-      theme="dark"
-      offset={{ top: 75, right: 16 }}
-      richColors
-    />
-  );
-};

frontend/realtime/src/types/ServerTypes.ts

@@ -0,0 +1,36 @@
+// Server communication types
+export interface ModelConfig {
+  name: string;
+  path: string;
+  precision: "fp16" | "fp32";
+}
+
+export interface LoraConfig {
+  path: string;
+  alpha: number;
+}
+
+export interface ServerSetupPayload {
+  type: "setup";
+  model: ModelConfig;
+  lora?: LoraConfig;
+  device: "cuda" | "cpu";
+}
+
+export interface ServerSettingsPayload {
+  type: "settings";
+  prompt?: string;
+  strength?: number;
+  steps?: number;
+  guidance_scale?: number;
+  seed?: number;
+  negative_prompt?: string;
+}
+
+export interface ServerResponse {
+  type: "progress" | "success" | "error";
+  status?: "model_loaded" | "generation_complete";
+  message?: string;
+  percent?: number;
+  image?: string; // base64 encoded image
+}

frontend/realtime/start_server.ps1

@@ -1,2 +1,5 @@
-.\.venv\Scripts\Activate.ps1
-python paint_py_server\paint_test.py
+# Activate virtual environment
+deactivate conda
+.\.venv\Scripts\activate.bat
+cd .\paint_py_server
+python server.py

frontend/realtime/yarn.lock

@@ -3635,7 +3635,7 @@ react-colorful@^5.6.1:
   resolved "https://registry.npmjs.org/react-colorful/-/react-colorful-5.6.1.tgz"
   integrity sha512-1exovf0uGTGyq5mXQT0zgQ80uvj2PCwvF8zY1RN9/vbJVSjSo3fsB/4L3ObbF7u70NduSiK4xu4Y6q1MHoUGEw==
 
-"react-dom@^16.8.0 || ^17.0.0 || ^18.0.0", react-dom@^18, react-dom@^18.0.0, "react-dom@^18.0.0 || ^19.0.0 || ^19.0.0-rc", react-dom@^18.3.1, "react-dom@>= 16.8.0", react-dom@>=16.4.0, react-dom@>=16.8, react-dom@>=16.8.0:
+"react-dom@^16.8.0 || ^17.0.0 || ^18.0.0", react-dom@^18, react-dom@^18.0.0, react-dom@^18.3.1, "react-dom@>= 16.8.0", react-dom@>=16.4.0, react-dom@>=16.8, react-dom@>=16.8.0:
   version "18.3.1"
   resolved "https://registry.npmjs.org/react-dom/-/react-dom-18.3.1.tgz"
   integrity sha512-5m4nQKp+rZRb09LNH59GM4BxTh9251/ylbKIbpe7TpGxfJ+9kv6BLkLBXIjjspbgbnIBNqlI23tRnTWT0snUIw==
@@ -3684,7 +3684,7 @@ react-router@6.26.1:
   dependencies:
     "@remix-run/router" "1.19.1"
 
-"react@^16.14.0 || 17.x || 18.x", "react@^16.8.0 || ^17.0.0 || ^18.0.0", "react@^16.8.0 || ^17.0.0-rc.1 || ^18.0.0 || ^19.0.0", react@^18, react@^18.0.0, "react@^18.0.0 || ^19.0.0 || ^19.0.0-rc", react@^18.3.1, "react@>= 16.8.0", react@>=16.3, react@>=16.4.0, react@>=16.8, react@>=16.8.0:
+"react@^16.14.0 || 17.x || 18.x", "react@^16.8.0 || ^17.0.0 || ^18.0.0", "react@^16.8.0 || ^17.0.0-rc.1 || ^18.0.0 || ^19.0.0", react@^18, react@^18.0.0, react@^18.3.1, "react@>= 16.8.0", react@>=16.3, react@>=16.4.0, react@>=16.8, react@>=16.8.0:
   version "18.3.1"
   resolved "https://registry.npmjs.org/react/-/react-18.3.1.tgz"
   integrity sha512-wS+hAgJShR0KhEvPJArfuPVN1+Hz1t0Y6n5jLrGQbkb4urgPE/0Rve+1kMB1v/oWgHgm4WIcV+i7F2pTVj+2iQ==
@@ -3956,11 +3956,6 @@ socks@^2.8.3:
     ip-address "^9.0.5"
     smart-buffer "^4.2.0"
 
-sonner@^2.0.1:
-  version "2.0.1"
-  resolved "https://registry.npmjs.org/sonner/-/sonner-2.0.1.tgz"
-  integrity sha512-FRBphaehZ5tLdLcQ8g2WOIRE+Y7BCfWi5Zyd8bCvBjiW8TxxAyoWZIxS661Yz6TGPqFQ4VLzOF89WEYhfynSFQ==
-
 source-map-js@^1.2.1:
   version "1.2.1"
   resolved "https://registry.npmjs.org/source-map-js/-/source-map-js-1.2.1.tgz"