toaster for notifications

swaps to shape selection or image selection when that is added,
except for render copy.

Cargo.lock

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

crates/desktop/src-tauri/Cargo.toml

@@ -48,3 +48,6 @@ 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,6 +1,15 @@
 // 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

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

frontend/realtime/package-lock.json

@@ -23,6 +23,7 @@
         "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"
@@ -8289,6 +8290,15 @@
         "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,6 +26,7 @@
     "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

@@ -0,0 +1,272 @@
+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

@@ -1,527 +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
-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

@@ -0,0 +1,101 @@
+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,8 +30,7 @@ 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;
@@ -64,6 +63,7 @@ export class Engine {
   private selectorSquare: SelectorSquare;
   private loadingVideosProvider: LoadingVideosProvider;
   private matteBox: MatteBox;
+  private backgroundNode: ShapeNode | null = null;
 
   private sceneManager: SceneManager;
   private undoStackManager: UndoStackManager;
@@ -196,19 +196,20 @@ 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();
@@ -216,6 +217,9 @@ 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;
@@ -310,66 +314,9 @@ 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(),
     );
@@ -439,39 +386,23 @@ export class Engine {
       this.addText(textdata);
     });
 
-    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,
-          });
-        }
-      }
+   
+    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.promptEvents.onPromptStrengthChanged((strength) => {
-      this.realTimeDrawEngine.currentStrength = strength;
-    });
-
-    uiEvents.promptEvents.onPromptTextChanged((prompt) => {
+    uiEvents.promptEvents.onPromptTextChanged(async (prompt) => {
       this.realTimeDrawEngine.currentPrompt = prompt;
+      clearTimeout(renderTimeout);
+      renderTimeout = setTimeout(async () => {
+        await this.realTimeDrawEngine.render();
+      }, 1000);
     });
 
     uiEvents.onAddShapeToEngine((shapeData) => {
@@ -488,6 +419,11 @@ 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) {
@@ -640,6 +576,7 @@ 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(),
@@ -650,9 +587,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);
 
@@ -667,6 +604,7 @@ export class Engine {
       selectionManagerRef: this.selectionManager,
       loaded: async () => {
         this.realTimeDrawEngine.render();
+        this.setAppMode(AppModes.SELECT);
       },
     });
     this.commandManager.createNode(copyNode);
@@ -680,9 +618,10 @@ 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,8 +31,7 @@ export class PaintNode extends NetworkedNode {
     selectionManagerRef,
     loaded,
   }: PaintNodeConstructor) {
-
-    const cl = canvasElement
+    const cl = canvasElement;
 
     // Convert circle to image
     const dataURL = cl.toDataURL();
@@ -40,31 +39,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,6 +180,8 @@ 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,14 +15,24 @@ import { Image } from "@tauri-apps/api/image";
 
 import { PaintNode } from "../Nodes/PaintNode";
 
-import {
-  ServerSetupPayload,
-  ServerSettingsPayload,
-  ServerResponse,
-} from "../types/ServerTypes";
+import { WebSocketClient } from "../../PyServer/DrawServer";
+import { PreviewCopyNode } from "../Nodes/PreviewCopy";
 
 // 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)[];
@@ -60,7 +70,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,
@@ -74,8 +84,12 @@ export class RealTimeDrawEngine {
 
   private onPreviewCopyCallback?: (previewCopy: Konva.Image) => void; // New Callback
 
-  private serverSocket: WebSocket | null = null;
+  private client: WebSocketClient | null = null;
+  private isConnected: boolean = false;
 
+
+  private backgroundNode: ShapeNode | null = null;
+  public backgroundColor: string = "#d2d2d2"
   constructor({
     width,
     height,
@@ -127,6 +141,7 @@ 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
@@ -136,7 +151,7 @@ export class RealTimeDrawEngine {
     this.fps = 24;
 
     this.currentPrompt = "";
-    this.currentStrength = 100;
+    this.currentStrength = 1;
     // This is captures a subset of the medialayer ref
     this.captureCanvas = new Konva.Rect({
       name: "CaptureCanvas",
@@ -144,7 +159,7 @@ export class RealTimeDrawEngine {
       y: this.positionY,
       width: this.width,
       height: this.height,
-      fill: "white",
+      fill: this.backgroundColor,
       stroke: "black",
       strokeWidth: 1,
       draggable: false,
@@ -171,102 +186,104 @@ 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 startServer() {
-    // Create WebSocket connection
-    const socket = new WebSocket("ws://localhost:8765");
+  public async startServer() {
+    try {
+      this.client = new WebSocketClient("ws://localhost:8765");
 
-    // 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
+      // Set up message handlers
+      this.client.onProgress((progress) => {
+        console.log(
+          `Progress: ${progress.message} (${progress.progress * 100}%)`,
+        );
+        if (progress.error) {
+          console.error("Error:", progress.error);
         }
-        // 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
+      });
+
+      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 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);
-      }
-    };
+      });
 
-    socket.onerror = (error) => {
-      console.error("WebSocket error:", error);
-    };
+      // Wait for connection to be established
+      await new Promise<void>((resolve, reject) => {
+        if (!this.client) return reject("Client not initialized");
 
-    socket.onclose = () => {
-      console.log("Disconnected from inference server");
-    };
+        this.client.ws.onopen = () => {
+          this.isConnected = true;
+          resolve();
+        };
+        this.client.ws.onerror = (error) => reject(error);
+      });
 
-    // Store socket reference for later use
-    this.serverSocket = socket;
+      // 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;
+    }
   }
 
-  // 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",
-      );
+  private async loadModel(settings: ServerSettings) {
+    if (!this.client || !this.isConnected) {
+      throw new Error("Server not connected");
     }
+    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();
@@ -287,7 +304,7 @@ export class RealTimeDrawEngine {
         y: this.captureCanvas.y(),
         draggable: false,
       });
-
+      this.drawingsLayer.moveToTop();
       this.drawingsLayer.add(currentLine); // Add to drawingsLayer
       isDrawing = true;
     };
@@ -360,9 +377,6 @@ 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)) {
@@ -392,17 +406,26 @@ 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();
   }
 
@@ -417,89 +440,96 @@ 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,
+    // 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(),
           });
-    
-          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();
-            }
-          });
-          
+
+          e.target.off("dragend");
+        } else {
           // Start Generation Here
-          const stage = this.mediaLayerRef.getStage();
-          const layer = this.mediaLayerRef; // Use the existing media layer
+          previewCopy.destroy();
+          this.mediaLayerRef.batchDraw();
+        }
+      });
 
-          // 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);
+      // Start Generation Here
+      const stage = this.mediaLayerRef.getStage();
+      const layer = this.mediaLayerRef; // Use the existing media layer
 
-          layer.add(centeredText);
-          centeredText.moveToTop();
+      // 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.batchDraw();
-    
-          // 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);
-        });
+        }
+        stage.off("mouseup touchend", removeText);
+      };
+      stage.on("mouseup touchend", removeText);
+    });
   }
 
   public findImageNodeById(
@@ -647,12 +677,7 @@ 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();
   };
 
@@ -661,19 +686,20 @@ export class RealTimeDrawEngine {
       node instanceof ImageNode ||
       node instanceof TextNode ||
       node instanceof ShapeNode ||
-      node instanceof PaintNode
+      node instanceof PaintNode ||
+      node instanceof PreviewCopyNode
     ) {
       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();
   }
 
@@ -688,7 +714,8 @@ export class RealTimeDrawEngine {
       node instanceof ImageNode ||
       node instanceof TextNode ||
       node instanceof ShapeNode ||
-      node instanceof PaintNode
+      node instanceof PaintNode ||
+      node instanceof PreviewCopyNode
     ) {
       const index = this.imageNodes.indexOf(node);
       if (index > -1) {
@@ -849,6 +876,12 @@ 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
@@ -872,9 +905,20 @@ export class RealTimeDrawEngine {
 
     // Test code
     if (true) {
-      this.outputBitmap = bitmap;
-      this.previewCanvas.image(bitmap);
-      this.isProcessing = false;
+      //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,
+      });
       return;
     }
 
@@ -903,10 +947,33 @@ export class RealTimeDrawEngine {
 
   // Add getter/setter for brush size
   public set paintBrushSize(size: number) {
-    this.brushSize = size;
+    this.brushSize =  size * 7.5;
+    const stage = this.mediaLayerRef.getStage();
+    if (stage && this.isEnabled) {
+      this.updateCursor(stage);
+    }
   }
 
   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/KonvaRootComponent/SignaledModelButtons.tsx

@@ -0,0 +1,45 @@
+import { twMerge } from "tailwind-merge";
+import { ModelButton } from "~/components/features/ModelSelection/ModelButton";
+import { useSignalEffect } from "@preact/signals-react";
+import { selectedModel, selectedLora } from "~/signals/uiEvents/modelSelection";
+import { getModelById, getLoraById } from "~/data/models";
+
+export const SignaledModelButtons = () => {
+  // This makes the component re-render when the selected model or lora changes
+  useSignalEffect(() => {
+    selectedModel.value;
+    selectedLora.value;
+  });
+
+  // Get the selected model's image or use default
+  const selectedModelData = selectedModel.value
+    ? getModelById(selectedModel.value)
+    : null;
+  const modelIconSrc = selectedModelData?.imageUrl || "";
+
+  // Get the selected LoRA's image or use default
+  const selectedLoraData = selectedLora.value
+    ? getLoraById(selectedLora.value)
+    : null;
+  const loraIconSrc = selectedLoraData?.imageUrl || "";
+
+  return (
+    <div className={twMerge("z-20 flex h-fit w-fit items-center gap-2")}>
+      <ModelButton
+        iconSrc={modelIconSrc}
+        iconAlt={selectedModelData?.name || "Model"}
+        label={selectedModelData?.name || "Realistic"}
+        subtitle="Model"
+        type="model"
+      />
+
+      <ModelButton
+        iconSrc={loraIconSrc}
+        iconAlt={selectedLoraData?.name || "LoRA"}
+        label={selectedLoraData?.name || "Detailed Tweaker"}
+        subtitle="LoRA"
+        type="lora"
+      />
+    </div>
+  );
+};

frontend/realtime/src/PyServer/DrawServer.ts

@@ -0,0 +1,67 @@
+// 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

@@ -0,0 +1,83 @@
+// 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

@@ -1,260 +0,0 @@
-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

@@ -1,216 +0,0 @@
-// 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,6 +5,7 @@ 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";
@@ -31,11 +32,17 @@ const App = () => {
       <StrictMode>
         <PostHogProvider client={posthog}>
           <RouterProvider router={router} />
+          <Toaster />
         </PostHogProvider>
       </StrictMode>
     );
   }
-  return <RouterProvider router={router} />;
+  return (
+    <>
+      <RouterProvider router={router} />
+      <Toaster />
+    </>
+  );
 };
 
 createRoot(document.getElementById("root")!).render(<App />);

frontend/realtime/src/components/features/ModelSelection/LoraSelectionDialog.tsx

@@ -0,0 +1,64 @@
+import { DialogTitle } from "@headlessui/react";
+import { useSignalEffect } from "@preact/signals-react";
+import { ModelCard } from "./ModelCard";
+import { loras, isLoraCompatibleWithModel } from "~/data/models";
+import {
+  selectedModel,
+  selectedLora,
+  dispatchers,
+} from "~/signals/uiEvents/modelSelection";
+
+interface LoraSelectionDialogProps {
+  onClose: () => void;
+}
+
+export const LoraSelectionDialog = ({ onClose }: LoraSelectionDialogProps) => {
+  // This makes component re-render whenever signals change
+  useSignalEffect(() => {
+    selectedModel.value;
+    selectedLora.value;
+  });
+
+  const handleSelectLora = (loraId: string) => {
+    dispatchers.setSelectedLora(loraId);
+    onClose(); // Close the dialog after selection
+  };
+
+  const currentModelId = selectedModel.value;
+
+  return (
+    <div className="flex flex-col gap-4">
+      <div>
+        <DialogTitle className="text-2xl font-bold">Select LoRA</DialogTitle>
+        <p className="text-sm text-gray-400">
+          Pick a LoRA to enhance your selected model
+          {!currentModelId && (
+            <span className="ml-1 text-yellow-500">
+              (Select a base model first)
+            </span>
+          )}
+        </p>
+      </div>
+
+      <div className="grid grid-cols-4 gap-4">
+        {loras.map((lora) => (
+          <ModelCard
+            key={lora.id}
+            id={lora.id}
+            name={lora.name}
+            imageUrl={lora.imageUrl}
+            isDownloaded={lora.isDownloaded}
+            isSelected={selectedLora.value === lora.id}
+            isCompatible={
+              currentModelId
+                ? isLoraCompatibleWithModel(lora.id, currentModelId)
+                : false
+            }
+            type="lora"
+            onSelect={handleSelectLora}
+          />
+        ))}
+      </div>
+    </div>
+  );
+};

frontend/realtime/src/components/features/ModelSelection/ModelButton.tsx

@@ -0,0 +1,93 @@
+import { useState } from "react";
+import { Button } from "~/components/ui/Button";
+import { BaseDialog } from "~/components/ui/BaseDialog";
+import { useSignalEffect } from "@preact/signals-react";
+import { selectedModel, selectedLora } from "~/signals/uiEvents/modelSelection";
+import { getModelById, getLoraById } from "~/data/models";
+import { ModelSelectionDialog } from "./ModelSelectionDialog";
+import { LoraSelectionDialog } from "./LoraSelectionDialog";
+
+interface ModelButtonProps {
+  iconSrc: string;
+  iconAlt: string;
+  label?: string;
+  subtitle?: string;
+  type: "model" | "lora";
+}
+
+export const ModelButton = ({
+  iconSrc,
+  iconAlt,
+  label,
+  subtitle,
+  type,
+}: ModelButtonProps) => {
+  const [isOpen, setIsOpen] = useState(false);
+
+  // Update button label based on selection
+  useSignalEffect(() => {
+    // This will re-render component when these signals change
+    selectedModel.value;
+    selectedLora.value;
+  });
+
+  // Set button text based on current selection
+  let buttonLabel = label;
+  let buttonSubtitle = subtitle;
+
+  if (type === "model" && selectedModel.value) {
+    const model = getModelById(selectedModel.value);
+    if (model) {
+      buttonLabel = model.name;
+      buttonSubtitle = "Model";
+    }
+  }
+
+  if (type === "lora" && selectedLora.value) {
+    const lora = getLoraById(selectedLora.value);
+    if (lora) {
+      buttonLabel = lora.name;
+      buttonSubtitle = "LoRA";
+    }
+  }
+
+  return (
+    <>
+      <Button
+        variant="secondary"
+        className="glass flex w-40 items-center justify-start gap-2 bg-ui-controls/70 p-1"
+        onClick={() => setIsOpen(true)}
+      >
+        <img
+          src={iconSrc}
+          alt={iconAlt}
+          className="h-8 w-8 rounded-md bg-white/10"
+        />
+        {buttonSubtitle ? (
+          <div className="flex flex-col gap-0 truncate text-start">
+            <span className="truncate text-[11px] opacity-60">
+              {buttonSubtitle}
+            </span>
+            <span className="truncate text-[13px] opacity-80">
+              {buttonLabel}
+            </span>
+          </div>
+        ) : (
+          <span className="truncate">{buttonLabel}</span>
+        )}
+      </Button>
+
+      <BaseDialog
+        isOpen={isOpen}
+        onClose={() => setIsOpen(false)}
+        className="max-w-3xl"
+      >
+        {type === "model" ? (
+          <ModelSelectionDialog onClose={() => setIsOpen(false)} />
+        ) : (
+          <LoraSelectionDialog onClose={() => setIsOpen(false)} />
+        )}
+      </BaseDialog>
+    </>
+  );
+};

frontend/realtime/src/components/features/ModelSelection/ModelCard.tsx

@@ -0,0 +1,156 @@
+import { faDownToLine } from "@fortawesome/pro-solid-svg-icons";
+import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
+import { twMerge } from "tailwind-merge";
+import { useSignalEffect } from "@preact/signals-react";
+import { useState } from "react";
+
+import {
+  downloadingItems,
+  dispatchers,
+} from "~/signals/uiEvents/modelSelection";
+import { downloadModel } from "~/services/modelDownloadService";
+
+interface ModelCardProps {
+  id: string;
+  name: string;
+  imageUrl: string;
+  isDownloaded: boolean;
+  isSelected: boolean;
+  isCompatible: boolean;
+  type: "model" | "lora";
+  onSelect: (id: string) => void;
+}
+
+export const ModelCard = ({
+  id,
+  name,
+  imageUrl,
+  isDownloaded,
+  isSelected,
+  isCompatible,
+  type,
+  onSelect,
+}: ModelCardProps) => {
+  // Local state to track download status and progress
+  const [progress, setProgress] = useState(0);
+  const [isDownloading, setIsDownloading] = useState(false);
+  const [localIsDownloaded, setLocalIsDownloaded] = useState(isDownloaded);
+
+  // Update local state when downloadingItems changes
+  useSignalEffect(() => {
+    const downloads = downloadingItems.value;
+    const isCurrentlyDownloading = !!downloads[id];
+    setIsDownloading(isCurrentlyDownloading);
+
+    if (isCurrentlyDownloading) {
+      setProgress(downloads[id].progress);
+    }
+  });
+
+  // Function to handle the download
+  const handleDownload = async () => {
+    try {
+      // Use the download service to start the download
+      await downloadModel(id, type);
+
+      // For the dummy implementation, we'll update the local state when download completes
+      const checkDownloadStatus = setInterval(() => {
+        if (!dispatchers.isDownloading(id)) {
+          clearInterval(checkDownloadStatus);
+          setLocalIsDownloaded(true);
+        }
+      }, 500);
+    } catch (error) {
+      console.error(`Error downloading ${type} ${id}:`, error);
+    }
+  };
+
+  const handleClick = () => {
+    if (!isCompatible) return; // Do nothing if not compatible
+
+    if (isDownloading) return; // Do nothing if already downloading
+
+    if (localIsDownloaded) {
+      onSelect(id);
+    } else {
+      handleDownload();
+    }
+  };
+
+  return (
+    <div
+      className={twMerge(
+        "relative flex aspect-square flex-col items-center overflow-hidden rounded-xl border-2 bg-white/5 transition-all",
+        isSelected ? "border-primary" : "border-transparent",
+        !isCompatible ? "opacity-50 grayscale" : "hover:border-primary",
+        "cursor-pointer",
+      )}
+      onClick={handleClick}
+    >
+      {/* Selected indicator */}
+      {isSelected && (
+        <div className="absolute right-2 top-2 z-10 flex h-6 w-6 items-center justify-center rounded-full bg-primary text-white">
+          <svg
+            xmlns="http://www.w3.org/2000/svg"
+            width="16"
+            height="16"
+            viewBox="0 0 24 24"
+            fill="none"
+            stroke="currentColor"
+            strokeWidth="3"
+            strokeLinecap="round"
+            strokeLinejoin="round"
+          >
+            <polyline points="20 6 9 17 4 12"></polyline>
+          </svg>
+        </div>
+      )}
+
+      {/* Image container */}
+      <div className="relative h-full w-full overflow-hidden">
+        {/* Downloading indicator */}
+        {isDownloading && (
+          <div className="absolute inset-0 flex h-full w-full flex-col items-center justify-center bg-gray-900/75">
+            <svg
+              className="h-8 w-8 animate-spin text-white/40"
+              viewBox="0 0 64 64"
+              fill="none"
+              xmlns="http://www.w3.org/2000/svg"
+              width="24"
+              height="24"
+            >
+              <path
+                d="M32 3C35.8083 3 39.5794 3.75011 43.0978 5.20749C46.6163 6.66488 49.8132 8.80101 52.5061 11.4939C55.199 14.1868 57.3351 17.3837 58.7925 20.9022C60.2499 24.4206 61 28.1917 61 32C61 35.8083 60.2499 39.5794 58.7925 43.0978C57.3351 46.6163 55.199 49.8132 52.5061 52.5061C49.8132 55.199 46.6163 57.3351 43.0978 58.7925C39.5794 60.2499 35.8083 61 32 61C28.1917 61 24.4206 60.2499 20.9022 58.7925C17.3837 57.3351 14.1868 55.199 11.4939 52.5061C8.801 49.8132 6.66487 46.6163 5.20749 43.0978C3.7501 39.5794 3 35.8083 3 32C3 28.1917 3.75011 24.4206 5.2075 20.9022C6.66489 17.3837 8.80101 14.1868 11.4939 11.4939C14.1868 8.80099 17.3838 6.66487 20.9022 5.20749C24.4206 3.7501 28.1917 3 32 3L32 3Z"
+                stroke="currentColor"
+                strokeWidth="5"
+                strokeLinecap="round"
+                strokeLinejoin="round"
+              ></path>
+              <path
+                d="M32 3C36.5778 3 41.0906 4.08374 45.1692 6.16256C49.2477 8.24138 52.7762 11.2562 55.466 14.9605C58.1558 18.6647 59.9304 22.9531 60.6448 27.4748C61.3591 31.9965 60.9928 36.6232 59.5759 40.9762"
+                stroke="currentColor"
+                strokeWidth="5"
+                strokeLinecap="round"
+                strokeLinejoin="round"
+                className="text-white/90"
+              ></path>
+            </svg>
+            <div className="mt-2 text-center font-medium text-white">
+              {progress}%
+            </div>
+          </div>
+        )}
+
+        <img src={imageUrl} alt={name} className="h-full w-full object-cover" />
+
+        {/* Title */}
+        <div className="absolute bottom-0 mt-2 flex w-full items-center justify-between bg-[#3E3E41]/70 px-2 py-1 text-sm font-medium backdrop-blur-[2px]">
+          {name}
+          {!localIsDownloaded && !isDownloading && (
+            <FontAwesomeIcon icon={faDownToLine} className="text-white" />
+          )}
+        </div>
+      </div>
+    </div>
+  );
+};

frontend/realtime/src/components/features/ModelSelection/ModelSelectionDialog.tsx

@@ -0,0 +1,56 @@
+import { DialogTitle } from "@headlessui/react";
+import { useSignalEffect } from "@preact/signals-react";
+import { ModelCard } from "./ModelCard";
+import { models } from "~/data/models";
+import { selectedModel, dispatchers } from "~/signals/uiEvents/modelSelection";
+import { useEffect } from "react";
+
+interface ModelSelectionDialogProps {
+  onClose: () => void;
+}
+
+export const ModelSelectionDialog = ({
+  onClose,
+}: ModelSelectionDialogProps) => {
+  // This makes component re-render whenever signals change
+  useSignalEffect(() => {
+    selectedModel.value;
+  });
+
+  // Ensure a model is selected when the dialog opens
+  useEffect(() => {
+    if (selectedModel.value === null && models.length > 0) {
+      dispatchers.setSelectedModel(models[0].id);
+    }
+  }, []);
+
+  const handleSelectModel = (modelId: string) => {
+    dispatchers.setSelectedModel(modelId);
+    onClose(); // Close the dialog after selection
+  };
+
+  return (
+    <div className="flex flex-col gap-4">
+      <div>
+        <DialogTitle className="text-2xl font-bold">Select Model</DialogTitle>
+        <p className="text-sm text-gray-400">Pick your base generation model</p>
+      </div>
+
+      <div className="grid grid-cols-4 gap-4">
+        {models.map((model) => (
+          <ModelCard
+            key={model.id}
+            id={model.id}
+            name={model.name}
+            imageUrl={model.imageUrl}
+            isDownloaded={model.isDownloaded}
+            isSelected={selectedModel.value === model.id}
+            isCompatible={true} // Base models are always compatible
+            type="model"
+            onSelect={handleSelectModel}
+          />
+        ))}
+      </div>
+    </div>
+  );
+};

frontend/realtime/src/components/features/ToolbarTopCenter/index.tsx

@@ -1,91 +1,11 @@
-import { twMerge } from "tailwind-merge";
-import { Button } from "~/components/ui/Button";
-import { useState } from "react";
-import { BaseDialog } from "~/components/ui/BaseDialog";
-import { DialogTitle } from "@headlessui/react";
 import { SignaledLoadingIndicator } from "~/KonvaRootComponent/SignaledLoadingIndicator";
-
-interface ModelButtonProps {
-  iconSrc: string;
-  iconAlt: string;
-  label?: string;
-  subtitle?: string;
-  children?: React.ReactNode;
-}
-
-const ModelButton = ({
-  iconSrc,
-  iconAlt,
-  label,
-  subtitle,
-  children,
-}: ModelButtonProps) => {
-  const [isOpen, setIsOpen] = useState(false);
-
-  return (
-    <>
-      <Button
-        variant="secondary"
-        className="glass flex w-40 items-center justify-start gap-2 bg-ui-controls/70 p-1"
-        onClick={() => setIsOpen(true)}
-      >
-        <img
-          src={iconSrc}
-          alt={iconAlt}
-          className="h-8 w-8 rounded-md bg-white/10"
-        />
-        {subtitle ? (
-          <div className="flex flex-col gap-0 truncate text-start">
-            <span className="truncate text-[11px] opacity-60">{subtitle}</span>
-            <span className="truncate text-[13px] opacity-80">{label}</span>
-          </div>
-        ) : (
-          <span className="truncate">{label}</span>
-        )}
-      </Button>
-
-      <BaseDialog isOpen={isOpen} onClose={() => setIsOpen(false)}>
-        {children}
-      </BaseDialog>
-    </>
-  );
-};
+import { SignaledModelButtons } from "~/KonvaRootComponent/SignaledModelButtons";
 
 export const ToolbarTopCenter = () => {
   return (
     <div className="flex flex-col gap-3">
       {/* Model selection buttons (only shown in dev mode for now) */}
-      {import.meta.env.DEV && (
-        <div className={twMerge("z-20 flex h-fit w-fit items-center gap-2")}>
-          <ModelButton
-            iconSrc="/resources/icons/model.png"
-            iconAlt="Model"
-            label="Realistic"
-            subtitle="Model"
-          >
-            <div className="flex flex-col gap-2 text-sm">
-              <DialogTitle className="text-2xl font-bold">
-                Select Model
-              </DialogTitle>
-              <p>Model configuration options will go here</p>
-            </div>
-          </ModelButton>
-
-          <ModelButton
-            iconSrc="/resources/icons/lora.png"
-            iconAlt="LoRA"
-            label="Detailed Tweaker"
-            subtitle="LoRA"
-          >
-            <div className="flex flex-col gap-2 text-sm">
-              <DialogTitle className="text-2xl font-bold">
-                Select LoRA
-              </DialogTitle>
-              <p>LoRA configuration options will go here</p>
-            </div>
-          </ModelButton>
-        </div>
-      )}
+      {import.meta.env.DEV && <SignaledModelButtons />}
 
       {/* Loading models indicator (only shown in dev mode for now) */}
       {import.meta.env.DEV && <SignaledLoadingIndicator />}

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

@@ -0,0 +1,12 @@
+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/data/models.ts

@@ -0,0 +1,143 @@
+// Model and LoRA type definitions
+export interface ModelData {
+  id: string;
+  name: string;
+  type: string;
+  imageUrl: string;
+  description?: string;
+  isDownloaded: boolean;
+  compatibleLoraIds?: string[]; // LoRAs that work with this model
+}
+
+export interface LoraData {
+  id: string;
+  name: string;
+  imageUrl: string;
+  description?: string;
+  isDownloaded: boolean;
+  compatibleModelIds: string[]; // Models this LoRA works with
+}
+
+// Available models dummy data
+export const models: ModelData[] = [
+  {
+    id: "realistic",
+    name: "Realistic",
+    type: "Realistic",
+    imageUrl: "/images/models/realistic.png",
+    description: "Photorealistic image generation",
+    isDownloaded: true,
+    compatibleLoraIds: [
+      "detailed_tweaker",
+      "scenery_enhancer",
+      "portrait_plus",
+    ],
+  },
+  {
+    id: "flat_2d",
+    name: "Flat 2D",
+    type: "Flat 2D",
+    imageUrl: "/images/models/flat_2d.png",
+    description: "Flat 2D illustration style",
+    isDownloaded: true,
+    compatibleLoraIds: ["toon_enhancer", "detailed_tweaker"],
+  },
+  {
+    id: "ghibli",
+    name: "Ghibli",
+    type: "Ghibli",
+    imageUrl: "/images/models/ghibli.png",
+    description: "Studio Ghibli inspired style",
+    isDownloaded: false,
+    compatibleLoraIds: ["anime_style"],
+  },
+  {
+    id: "scenery",
+    name: "Scenery",
+    type: "Scenery",
+    imageUrl: "/images/models/scenery.png",
+    description: "Beautiful landscape generation",
+    isDownloaded: true,
+    compatibleLoraIds: ["scenery_enhancer"],
+  },
+  {
+    id: "concept",
+    name: "Concept",
+    type: "Concept",
+    imageUrl: "/images/models/concept.png",
+    description: "Abstract concept art generation",
+    isDownloaded: false,
+    compatibleLoraIds: ["detailed_tweaker", "scenery_enhancer"],
+  },
+];
+
+// Available LoRAs
+export const loras: LoraData[] = [
+  {
+    id: "detailed_tweaker",
+    name: "Detailed Tweaker",
+    imageUrl: "/resources/loras/detailedtweaker.jpg",
+    description: "Enhances fine details in generated images",
+    isDownloaded: true,
+    compatibleModelIds: [
+      "realistic",
+      "concept",
+      "anime",
+      "flat_2d",
+      "flux_realtime",
+    ],
+  },
+  {
+    id: "portrait_plus",
+    name: "Portrait+",
+    imageUrl: "/resources/loras/portraitplus.jpg",
+    description: "Improves portrait generation",
+    isDownloaded: false,
+    compatibleModelIds: ["realistic", "flux_realtime"],
+  },
+  {
+    id: "scenery_enhancer",
+    name: "Scenery Enhancer",
+    imageUrl: "/resources/loras/sceneryenhancer.jpg",
+    description: "Enhances landscape generation",
+    isDownloaded: true,
+    compatibleModelIds: ["realistic", "concept", "scenery"],
+  },
+  {
+    id: "anime_style",
+    name: "Anime Style+",
+    imageUrl: "/resources/loras/animestyle.jpg",
+    description: "Enhances anime art style",
+    isDownloaded: false,
+    compatibleModelIds: ["anime", "ghibli"],
+  },
+  {
+    id: "toon_enhancer",
+    name: "Toon Enhancer",
+    imageUrl: "/resources/loras/toonenhancer.jpg",
+    description: "Enhances cartoon styles",
+    isDownloaded: true,
+    compatibleModelIds: ["cartoon", "flat_2d"],
+  },
+];
+
+// Helper functions
+export const getModelById = (id: string): ModelData | undefined => {
+  return models.find((model) => model.id === id);
+};
+
+export const getLoraById = (id: string): LoraData | undefined => {
+  return loras.find((lora) => lora.id === id);
+};
+
+export const isLoraCompatibleWithModel = (
+  loraId: string,
+  modelId: string,
+): boolean => {
+  const lora = getLoraById(loraId);
+  return lora ? lora.compatibleModelIds.includes(modelId) : false;
+};
+
+export const getCompatibleLorasForModel = (modelId: string): LoraData[] => {
+  return loras.filter((lora) => lora.compatibleModelIds.includes(modelId));
+};

frontend/realtime/src/services/modelDownloadService.ts

@@ -0,0 +1,116 @@
+import { dispatchers } from "~/signals/uiEvents/modelSelection";
+import { models, loras, ModelData, LoraData } from "~/data/models";
+
+//THIS IS A VERY DUMMY DOWNLOAD IMPLEMENTATION AND SHOULD BE REPLACED WITH A REAL ONE FROM THE BACKEND OR SOMETHING BECAUSE IM NOT SURE - BFlat
+
+// // Interface for download progress events from Tauri backend
+// interface DownloadProgressEvent {
+//   id: string;
+//   progress: number;
+//   type: "model" | "lora";
+// }
+
+// // Interface for download completion events from Tauri backend
+// interface DownloadCompleteEvent {
+//   id: string;
+//   type: "model" | "lora";
+//   success: boolean;
+//   error?: string;
+// }
+
+// Initialize listeners for download events (dummy implementation)
+export const initializeDownloadListeners = async () => {
+  console.log("Download listeners initialized (dummy implementation)");
+  // No actual listeners in the dummy implementation
+  return Promise.resolve();
+};
+
+// Start downloading a model or lora (dummy implementation)
+export const downloadModel = async (id: string, type: "model" | "lora") => {
+  try {
+    // Start tracking the download in our state
+    dispatchers.startDownload(id, type);
+
+    // Simulate download progress
+    simulateDownloadProgress(id, type);
+
+    return true;
+  } catch (error) {
+    console.error(`Failed to start download for ${type} ${id}:`, error);
+    dispatchers.finishDownload(id);
+    return false;
+  }
+};
+
+// Simulate download progress
+const simulateDownloadProgress = (id: string, type: "model" | "lora") => {
+  let progress = 0;
+
+  // Create an interval to update progress
+  const interval = setInterval(() => {
+    progress += 5;
+
+    // Update progress in state
+    dispatchers.updateDownloadProgress(id, progress);
+
+    // When download completes
+    if (progress >= 100) {
+      clearInterval(interval);
+
+      // Finish the download
+      setTimeout(() => {
+        dispatchers.finishDownload(id);
+
+        // Update the local data to mark the model/lora as downloaded
+        updateDownloadedState(id, type, true);
+
+        console.log(`Download completed for ${type} ${id}`);
+      }, 500);
+    }
+  }, 200); // Update every 200ms
+};
+
+// Update the downloaded state in our local data
+const updateDownloadedState = (
+  id: string,
+  type: "model" | "lora",
+  isDownloaded: boolean,
+) => {
+  if (type === "model") {
+    // Find and update the model
+    const modelIndex = models.findIndex((model) => model.id === id);
+    if (modelIndex >= 0) {
+      const updatedModel: ModelData = {
+        ...models[modelIndex],
+        isDownloaded,
+      };
+
+      // Update the models array
+      models[modelIndex] = updatedModel;
+    }
+  } else {
+    // Find and update the lora
+    const loraIndex = loras.findIndex((lora) => lora.id === id);
+    if (loraIndex >= 0) {
+      const updatedLora: LoraData = {
+        ...loras[loraIndex],
+        isDownloaded,
+      };
+
+      // Update the loras array
+      loras[loraIndex] = updatedLora;
+    }
+  }
+};
+
+// Get all available models (dummy implementation)
+export const fetchAvailableModels = async (): Promise<ModelData[]> => {
+  // Just return the models from the data file
+  return Promise.resolve([...models]);
+};
+
+// Get all available loras (dummy implementation)
+export const fetchAvailableLoras = async (): Promise<LoraData[]> => {
+  // Just return the loras from the data file
+  return Promise.resolve([...loras]);
+};

frontend/realtime/src/signals/uiEvents/index.ts

@@ -28,6 +28,14 @@ import {
   dispatchers as promptDispatchers,
   events as promptEvents,
 } from "./promptSettings";
+import {
+  dispatchers as loadingIndicatorDispatchers,
+  events as loadingIndicatorEvents,
+} from "./loadingIndicator";
+import {
+  dispatchers as modelSelectionDispatchers,
+  events as modelSelectionEvents,
+} from "./modelSelection";
 
 export const uiEvents = {
   ...addMediaEvents,
@@ -39,6 +47,8 @@ export const uiEvents = {
   toolbarNode: toolbarNodeEvents,
   toolbarVideoExtraction: toolbarVideoExtractionEvents,
   promptEvents: promptEvents,
+  loadingIndicator: loadingIndicatorEvents,
+  modelSelection: modelSelectionEvents,
 };
 export const dispatchUiEvents = {
   ...addMediaDispatchers,
@@ -50,4 +60,6 @@ export const dispatchUiEvents = {
   toolbarNode: toolbarNodeDispatchers,
   toolbarVideoExtraction: toolbarVideoExtractionDispatchers,
   promptSettings: promptDispatchers,
+  loadingIndicator: loadingIndicatorDispatchers,
+  modelSelection: modelSelectionDispatchers,
 };

frontend/realtime/src/signals/uiEvents/modelSelection.ts

@@ -0,0 +1,100 @@
+import { effect, signal } from "@preact/signals-react";
+import { models } from "~/data/models";
+
+// Define the signals for selected models and download states
+export const selectedModel = signal<string | null>(
+  models.length > 0 ? models[0].id : null,
+);
+export const selectedLora = signal<string | null>(null);
+
+// Replace single download tracking with a map to track multiple downloads
+export interface DownloadInfo {
+  progress: number;
+  type: "model" | "lora";
+}
+
+// Map of modelId/loraId to download progress
+export const downloadingItems = signal<Record<string, DownloadInfo>>({});
+
+// Dispatchers
+const setSelectedModel = (modelId: string | null) => {
+  selectedModel.value = modelId;
+};
+
+const setSelectedLora = (loraId: string | null) => {
+  selectedLora.value = loraId;
+};
+
+// Start downloading a model or lora
+const startDownload = (id: string, type: "model" | "lora") => {
+  const newDownloads = { ...downloadingItems.value };
+  newDownloads[id] = { progress: 0, type };
+  downloadingItems.value = newDownloads;
+};
+
+// Update the progress of a downloading item
+const updateDownloadProgress = (id: string, progress: number) => {
+  if (!downloadingItems.value[id]) return;
+
+  const newDownloads = { ...downloadingItems.value };
+  newDownloads[id] = {
+    ...newDownloads[id],
+    progress: Math.max(0, Math.min(100, Math.round(progress))),
+  };
+  downloadingItems.value = newDownloads;
+};
+
+// Complete or cancel a download
+const finishDownload = (id: string) => {
+  const newDownloads = { ...downloadingItems.value };
+  delete newDownloads[id];
+  downloadingItems.value = newDownloads;
+};
+
+// Check if an item is currently downloading
+const isDownloading = (id: string): boolean => {
+  return !!downloadingItems.value[id];
+};
+
+// Get the progress of a downloading item
+const getDownloadProgress = (id: string): number => {
+  return downloadingItems.value[id]?.progress || 0;
+};
+
+// Event listeners
+const onSelectedModelChanged = (callback: (modelId: string | null) => void) => {
+  effect(() => {
+    callback(selectedModel.value);
+  });
+};
+
+const onSelectedLoraChanged = (callback: (loraId: string | null) => void) => {
+  effect(() => {
+    callback(selectedLora.value);
+  });
+};
+
+const onDownloadsChanged = (
+  callback: (downloads: Record<string, DownloadInfo>) => void,
+) => {
+  effect(() => {
+    callback(downloadingItems.value);
+  });
+};
+
+// EXPORTS
+export const dispatchers = {
+  setSelectedModel,
+  setSelectedLora,
+  startDownload,
+  updateDownloadProgress,
+  finishDownload,
+  isDownloading,
+  getDownloadProgress,
+};
+
+export const events = {
+  onSelectedModelChanged,
+  onSelectedLoraChanged,
+  onDownloadsChanged,
+};

frontend/realtime/src/types/ServerTypes.ts

@@ -1,36 +0,0 @@
-// 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,5 +1,2 @@
-# Activate virtual environment
-deactivate conda
-.\.venv\Scripts\activate.bat
-cd .\paint_py_server
-python server.py
+.\.venv\Scripts\Activate.ps1
+python paint_py_server\paint_test.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.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.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:
   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.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.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:
   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,6 +3956,11 @@ 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"