render bug fix

- Both modes' menus have been separated for better maintenance
- Slider is generic, can be reused
- New signals added for engine to hook into
- Sliders tested working on the UI side

frontend/realtime/.gitignore

@@ -12,7 +12,7 @@ dist
 dist-ssr
 *.local
 backup
-
+.venv
 # Editor directories and files
 ~
 .vscode

frontend/realtime/paint_py_server/server.py

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

frontend/realtime/src/KonvaApp/Engine.ts

@@ -2,7 +2,8 @@ import Konva from "konva";
 
 import { uiAccess, uiEvents } from "~/signals";
 import { ShapeNode } from "./Nodes";
-
+import { PaintNode } from "./Nodes/PaintNode";
+import { PreviewCopyNode } from "./Nodes/PreviewCopy";
 import { UndoStackManager } from "./UndoRedo";
 import { CommandManager, MatteBox, SceneManager } from "./EngineUtitlities";
 import {
@@ -110,17 +111,8 @@ export class Engine {
     });
 
     // core layer for all the work done.
-
     this.offScreenCanvas = new OffscreenCanvas(0, 0);
 
-    this.realTimeDrawEngine = new RealTimeDrawEngine({
-      width: VideoResolutions.SQUARE_1024.width,
-      height: VideoResolutions.SQUARE_1024.height,
-      mediaLayerRef: this.mediaLayer,
-      bgLayerRef: this.bgLayer,
-      offScreenCanvas: this.offScreenCanvas,
-    });
-
     // Collection of all Nodes
     this.nodesManager = new NodesManager();
     // Partial Collection of selected Nodes
@@ -128,6 +120,20 @@ export class Engine {
       nodeTransformerRef: this.nodeTransformer,
       mediaLayerRef: this.mediaLayer,
     });
+
+    this.realTimeDrawEngine = new RealTimeDrawEngine({
+      width: VideoResolutions.SQUARE_1024.width,
+      height: VideoResolutions.SQUARE_1024.height,
+      mediaLayerRef: this.mediaLayer,
+      offScreenCanvas: this.offScreenCanvas,
+      onDraw: async (canvas, lineBounds) => {
+        await this.addPaintNode(canvas, lineBounds);
+      },
+      onPreviewCopy: async (image) => {
+        await this.addPreviewCopy(image);
+      },
+    });
+
     // Selector Square to select Nodes
     this.selectorSquare = new SelectorSquare({
       captureCanvasRef: this.realTimeDrawEngine.captureCanvas,
@@ -139,7 +145,9 @@ export class Engine {
     this.uiLayer.add(this.selectorSquare.getKonvaNode());
 
     //Collection of commands for undo-redo
-    this.undoStackManager = new UndoStackManager();
+    this.undoStackManager = new UndoStackManager(() => {
+      this.realTimeDrawEngine.render();
+    });
     this.commandManager = new CommandManager({
       mediaLayerRef: this.mediaLayer,
       nodesManagerRef: this.nodesManager,
@@ -177,14 +185,31 @@ export class Engine {
     // some of the managers has events
     // hence, lastly, setup these events
     this.setupEventSystem();
-    this.setAppMode(AppModes.SELECT);
+    this.setAppMode(this.appMode);
   }
 
   private setAppMode(newAppMode: AppModes) {
     this.appMode = newAppMode;
     switch (this.appMode) {
+      case AppModes.PAINT: {
+        console.log("APPMODE: PAINT");
+        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,
+        });
+        this.matteBox.disable();
+        return;
+      }
       case AppModes.SELECT: {
         console.log("APPMODE: SELECT");
+        this.realTimeDrawEngine.disablePaintMode();
+        this.realTimeDrawEngine.enableDragging();
+
         this.selectorSquare.enable();
         this.selectionManager.enable();
         uiAccess.toolbarMain.enable();
@@ -356,8 +381,13 @@ export class Engine {
     );
 
     // Listen to Toolbar Main
-    uiEvents.toolbarMain.UNDO.onClick(() => this.undoStackManager.undo());
-    uiEvents.toolbarMain.REDO.onClick(() => this.undoStackManager.redo());
+    uiEvents.toolbarMain.UNDO.onClick(() => {
+      this.undoStackManager.undo();
+    });
+    uiEvents.toolbarMain.REDO.onClick(() => {
+      this.undoStackManager.redo();
+      this.realTimeDrawEngine.render();
+    });
 
     uiEvents.toolbarMain.SAVE.onClick(async (/*event*/) => {
       await this.realTimeDrawEngine.saveOutput();
@@ -365,10 +395,12 @@ export class Engine {
 
     uiEvents.toolbarMain.SELECT.onClick(() => {
       console.log("Toolbar Main >> Select");
+      this.setAppMode(AppModes.SELECT);
     });
 
     uiEvents.toolbarMain.PAINT.onClick(() => {
       console.log("Toolbar Main >> Paint");
+      this.setAppMode(AppModes.PAINT);
     });
 
     uiEvents.toolbarMain.ERASER.onClick(() => {
@@ -468,6 +500,9 @@ export class Engine {
       }
       this.changeNodeColor(nodeColor);
     });
+    uiEvents.toolbarMain.onPaintColorChanged((color) => {
+      this.realTimeDrawEngine.paintColor = color;
+    });
   }
 
   disableAllButtons() {
@@ -575,10 +610,13 @@ export class Engine {
           color: nodeColor.color,
           mediaLayerRef: this.mediaLayer,
           selectionManagerRef: this.selectionManager,
+          loaded: async () => {
+            this.realTimeDrawEngine.render();
+          },
         });
         newShapeNode.kNode.position(shapeNode.kNode.position());
         newShapeNode.kNode.zIndex(shapeNode.kNode.zIndex());
-        
+
         // Remove old node and add new one
         shapeNode.kNode.destroy();
         this.commandManager.createNode(newShapeNode);
@@ -608,19 +646,27 @@ export class Engine {
       selectionManagerRef: this.selectionManager,
       loaded: async () => {
         await this.realTimeDrawEngine.render();
-    },
+      },
     });
     shapeNode.kNode.zIndex(1); // Replace desiredZIndex with a number
 
     this.commandManager.createNode(shapeNode);
 
-
     console.debug("Added shapenode:", shapeNode);
     console.debug("Added node's ID:", shapeNode.kNode._id);
-
-    this.realTimeDrawEngine.addNodes(shapeNode);
   }
-
+  public addPreviewCopy(image: Image.Konva) {
+    // Start of Selection
+    const copyNode = new PreviewCopyNode({
+      image: image,
+      mediaLayerRef: this.mediaLayer,
+      selectionManagerRef: this.selectionManager,
+      loaded: async () => {
+        this.realTimeDrawEngine.render();
+      },
+    });
+    this.commandManager.createNode(copyNode);
+  }
   public addImage(imageFile: File) {
     const imageNode = new ImageNode({
       mediaLayerRef: this.mediaLayer,
@@ -634,8 +680,26 @@ export class Engine {
     });
     imageNode.kNode.zIndex(1);
     this.commandManager.createNode(imageNode);
-
-    this.realTimeDrawEngine.addNodes(imageNode);
+  }
+  public addPaintNode(
+    canvas: HTMLCanvasElement,
+    lineBounds: {
+      width: number;
+      height: number;
+      x: number;
+      y: number;
+    },
+  ) {
+    var node = new PaintNode({
+      canvasElement: canvas,
+      lineBounds: lineBounds,
+      mediaLayerRef: this.mediaLayer,
+      selectionManagerRef: this.selectionManager,
+      loaded: async () => {
+        await this.realTimeDrawEngine.render();
+      },
+    });
+    this.commandManager.createNode(node);
   }
 
   public addVideo(

frontend/realtime/src/KonvaApp/EngineUtitlities/CommandManager.ts

@@ -20,14 +20,15 @@ import {
   NodeTransformer,
   SelectionManager,
 } from "../NodesManagers";
-import { RenderEngine } from "../RenderingPrimitives/RenderEngine";
+import { RealTimeDrawEngine } from "../RenderingPrimitives/RealTimeDrawEngine";
+import { PaintNode } from "../Nodes/PaintNode";
 
 interface EngineReferences {
   mediaLayerRef: Konva.Layer;
   nodesManagerRef: NodesManager;
   nodeTransformerRef: NodeTransformer;
   selectionManagerRef: SelectionManager;
-  renderEngineRef: RenderEngine;
+  renderEngineRef: RealTimeDrawEngine;
   undoStackManagerRef: UndoStackManager;
 }
 
@@ -36,7 +37,7 @@ export class CommandManager {
   private nodesManagerRef: NodesManager;
   private nodeTransformerRef: NodeTransformer;
   private selectionManagerRef: SelectionManager;
-  private renderEngineRef: RenderEngine;
+  private renderEngineRef: RealTimeDrawEngine;
   private undoStackManagerRef: UndoStackManager;
 
   constructor(engineRefs: EngineReferences) {
@@ -48,7 +49,7 @@ export class CommandManager {
     this.undoStackManagerRef = engineRefs.undoStackManagerRef;
   }
 
-  createNode(node: VideoNode | ImageNode | TextNode | ShapeNode) {
+  createNode(node: MediaNode) {
     const command = new CreateCommand({
       nodes: new Set<MediaNode>([node]),
       mediaLayerRef: this.mediaLayerRef,
@@ -62,6 +63,7 @@ export class CommandManager {
     // Set the kNode's manual id to same as kNode's internal _id
     // This helps identify the kNode in the Konva layer with find
     node.kNode.id(node.kNode._id.toString());
+    this.renderEngineRef.render();
   }
   deleteNodes() {
     const nodes = this.selectionManagerRef.getSelectedNodes();
@@ -74,6 +76,7 @@ export class CommandManager {
       renderEngineRef: this.renderEngineRef,
     });
     this.undoStackManagerRef.executeCommand(command);
+    this.renderEngineRef.render();
   }
   toggleLockNodes() {
     const nodes = this.selectionManagerRef.getSelectedNodes();
@@ -101,6 +104,7 @@ export class CommandManager {
       mediaLayerRef: this.mediaLayerRef,
     });
     this.undoStackManagerRef.executeCommand(command);
+    this.renderEngineRef.render();
   }
   moveNodesDown() {
     const command = new MoveLayerDown({
@@ -109,6 +113,7 @@ export class CommandManager {
       mediaLayerRef: this.mediaLayerRef,
     });
     this.undoStackManagerRef.executeCommand(command);
+    this.renderEngineRef.render();
   }
   translateNodes(props: {
     nodes: Set<MediaNode>;
@@ -120,6 +125,7 @@ export class CommandManager {
       layerRef: this.mediaLayerRef,
     });
     this.undoStackManagerRef.pushCommand(command);
+    this.renderEngineRef.render();
   }
   transformNodes(props: {
     nodes: Set<MediaNode>;
@@ -131,6 +137,7 @@ export class CommandManager {
       layerRef: this.mediaLayerRef,
     });
     this.undoStackManagerRef.pushCommand(command);
+    this.renderEngineRef.render();
   }
   addChromaKey(props: { videoNode: VideoNode; newChromaColor: RGBColor }) {
     const command = new AddChromaKey(props);

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

@@ -0,0 +1,92 @@
+import Konva from "konva";
+import { SelectionManager } from "../NodesManagers";
+import { Position, NodeData } from "../types";
+import { NodeType } from "./constants";
+
+import { NetworkedNode } from "./NetworkedNode";
+
+interface PaintNodeConstructor {
+  canvasElement:HTMLCanvasElement;
+  lineBounds: {
+    width: number;
+    height: number;
+    x: number;
+    y: number;
+    };
+  mediaLayerRef: Konva.Layer;
+  selectionManagerRef: SelectionManager;
+  
+  loaded: () => Promise<void>;
+}
+
+export class PaintNode extends NetworkedNode {
+  retry(): void {
+    throw new Error("Method not implemented.");
+  }
+
+  constructor({
+    canvasElement,
+    lineBounds,
+    mediaLayerRef,
+    selectionManagerRef,
+    loaded,
+  }: PaintNodeConstructor) {
+
+    const cl = canvasElement
+
+    // Convert circle to image
+    const dataURL = cl.toDataURL();
+    const imageObj = new Image();
+    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' 
+      });
+
+    imageObj.onload = () => {
+        this.mediaLayerRef.draw();
+        loaded();
+    };
+  
+    super({
+      selectionManagerRef: selectionManagerRef,
+      mediaLayerRef: mediaLayerRef,
+      kNode: paintShape,
+    });
+
+    this.mediaLayerRef.add(this.kNode);
+    this.listenToBaseKNode();
+    this.mediaLayerRef.draw();
+  }
+
+  public getNodeData(canvasPosition: Position) {
+    const data: NodeData = {
+      type: NodeType.SHAPE,
+      transform: {
+        position: {
+          x: this.kNode.position().x - canvasPosition.x,
+          y: this.kNode.position().y - canvasPosition.y,
+        },
+        size: this.kNode.size(),
+        rotation: this.kNode.rotation(),
+        scale: {
+          x: this.kNode.scaleX(),
+          y: this.kNode.scaleY(),
+        },
+        zIndex: this.kNode.getZIndex(),
+      },
+    };
+    return data;
+  }
+}

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

@@ -0,0 +1,58 @@
+import Konva from "konva";
+import { SelectionManager } from "../NodesManagers";
+import { Position, Size, NodeData, TransformationData } from "../types";
+import { NodeType } from "./constants";
+import { NodeUtilities } from "./NodeUtilities";
+import { NetworkedNode } from "./NetworkedNode";
+
+interface PreviewCopyConstructor {
+  image: Konva.Image;
+  mediaLayerRef: Konva.Layer;
+  selectionManagerRef: SelectionManager;
+  loaded: () => Promise<void>;
+}
+
+export class PreviewCopyNode extends NetworkedNode {
+  retry(): void {
+    throw new Error("Method not implemented.");
+  }
+
+  constructor({
+    image,
+    mediaLayerRef,
+    selectionManagerRef,
+    loaded,
+  }: PreviewCopyConstructor) {
+    // Create the actual shape inside the group
+
+    super({
+      selectionManagerRef: selectionManagerRef,
+      mediaLayerRef: mediaLayerRef,
+      kNode: image,
+    });
+
+    this.mediaLayerRef.add(this.kNode);
+    this.listenToBaseKNode();
+    this.mediaLayerRef.draw();
+  }
+
+  public getNodeData(canvasPosition: Position) {
+    const data: NodeData = {
+      type: NodeType.SHAPE,
+      transform: {
+        position: {
+          x: this.kNode.position().x - canvasPosition.x,
+          y: this.kNode.position().y - canvasPosition.y,
+        },
+        size: this.kNode.size(),
+        rotation: this.kNode.rotation(),
+        scale: {
+          x: this.kNode.scaleX(),
+          y: this.kNode.scaleY(),
+        },
+        zIndex: this.kNode.getZIndex(),
+      },
+    };
+    return data;
+  }
+}

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

@@ -11,12 +11,21 @@ import { MediaNode } from "../types";
 
 import { RenderTask } from "./RenderTask";
 import { OffScreenSceneCanvas } from "./OffScreenSceneCanvas";
+import { Image } from "@tauri-apps/api/image";
+
+import { PaintNode } from "../Nodes/PaintNode";
+
+import {
+  ServerSetupPayload,
+  ServerSettingsPayload,
+  ServerResponse,
+} from "../types/ServerTypes";
 
 // https://www.aiseesoft.com/resource/phone-aspect-ratio-screen-resolution.html#:~:text=16%3A9%20Aspect%20Ratio
 
 export class RealTimeDrawEngine {
   private videoNodes: VideoNode[];
-  private imageNodes: (ImageNode | TextNode | ShapeNode)[];
+  private imageNodes: (ImageNode | TextNode | ShapeNode | PaintNode)[];
 
   private offScreenCanvas: OffscreenCanvas;
   private outputBitmap: ImageBitmap | undefined;
@@ -26,6 +35,7 @@ export class RealTimeDrawEngine {
   // capturing composite within window
 
   private mediaLayerRef: Konva.Layer;
+  private drawingsLayer: Konva.Layer; // New Layer for Drawings
 
   private height: number;
   private width: number;
@@ -46,20 +56,55 @@ export class RealTimeDrawEngine {
   public currentPrompt: string;
   public currentStrength: number;
 
+  // Paint Color
+  // paint Brush Size
+  // has to exit out of paint mode when shape or image are used.
+  public paintColor: string = "#000000";
+  public paintBrushSize: number = 5;
+  public isPaintMode: boolean = false;
+
+  private onDrawCallback?: (
+    canvas: HTMLCanvasElement,
+    lineBounds: {
+      width: number;
+      height: number;
+      x: number;
+      y: number;
+    },
+  ) => void;
+
+  private onPreviewCopyCallback?: (previewCopy: Konva.Image) => void; // New Callback
+
+  private serverSocket: WebSocket | null = null;
+
   constructor({
     width,
     height,
     mediaLayerRef,
     offScreenCanvas,
+    onDraw,
+    onPreviewCopy, // New Parameter
   }: {
     width: number;
     height: number;
     mediaLayerRef: Konva.Layer;
     offScreenCanvas: OffscreenCanvas;
+    onDraw?: (
+      canvas: HTMLCanvasElement,
+      lineBounds: {
+        width: number;
+        height: number;
+        x: number;
+        y: number;
+      },
+    ) => void;
+    onPreviewCopy?: (previewCopy: Konva.Image) => void; // New Parameter
   }) {
     this.videoLoadingCanvas = undefined;
     this.videoNodes = [];
     this.imageNodes = [];
+    this.onDrawCallback = onDraw;
+    this.onPreviewCopyCallback = onPreviewCopy; // Assign Callback
 
     // TODO: Make this dynamic and update this on change of canvas.
 
@@ -79,6 +124,12 @@ export class RealTimeDrawEngine {
     // this is the whole canvas
     this.mediaLayerRef = mediaLayerRef;
 
+    // Create a separate layer for drawings
+    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
 
     this.port = undefined;
@@ -121,24 +172,264 @@ export class RealTimeDrawEngine {
     // Add mouse events for preview canvas copying
     //this.previewCopyListener();
 
+    //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");
+
+    // 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
+        }
+        // Handle successful model load
+        else if (
+          response.type === "success" &&
+          response.status === "model_loaded"
+        ) {
+          console.log("Model loaded successfully");
+          // Enable UI elements that depend on model being loaded
+        }
+        // Handle image generation results
+        else if (
+          response.type === "success" &&
+          response.status === "generation_complete"
+        ) {
+          if (response.image) {
+            // Convert base64 image to ImageBitmap for preview
+            this.base64ToImageBitmap(response.image).then((bitmap) => {
+              this.outputBitmap = bitmap;
+              this.previewCanvas.image(bitmap);
+              this.mediaLayerRef.draw();
+            });
+          }
+        }
+      } catch (error) {
+        console.error("Error parsing server response:", error);
+      }
+    };
+
+    socket.onerror = (error) => {
+      console.error("WebSocket error:", error);
+    };
+
+    socket.onclose = () => {
+      console.log("Disconnected from inference server");
+    };
+
+    // Store socket reference for later use
+    this.serverSocket = socket;
+  }
+
+  // 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",
+      );
+    }
+  }
+
+  public paintMode() {
+    let isDrawing = false;
+    let currentLine: Konva.Line | null = null;
+
+    const startDrawing = (pos: { x: number; y: number }) => {
+      if (!this.isEnabled) return;
+      const stage = this.mediaLayerRef.getStage();
+      if (!stage) return;
+
+      // Convert pointer position to relative position within capture canvas
+      const captureBox = this.captureCanvas.getClientRect();
+      const relativeX = pos.x - captureBox.x;
+      const relativeY = pos.y - captureBox.y;
+
+      currentLine = new Konva.Line({
+        points: [relativeX, relativeY],
+        stroke: this.paintColor,
+        strokeWidth: 5,
+        lineCap: "round",
+        lineJoin: "round",
+        x: this.captureCanvas.x(),
+        y: this.captureCanvas.y(),
+        draggable: false,
+      });
+
+      this.drawingsLayer.add(currentLine); // Add to drawingsLayer
+      isDrawing = true;
+    };
+
+    const draw = (pos: { x: number; y: number }) => {
+      if (!this.isEnabled) return;
+      if (!isDrawing || !currentLine) return;
+
+      // Convert pointer position to relative position
+      const captureBox = this.captureCanvas.getClientRect();
+      const relativeX = pos.x - captureBox.x;
+      const relativeY = pos.y - captureBox.y;
+
+      const newPoints = currentLine.points().concat([relativeX, relativeY]);
+      currentLine.points(newPoints);
+      this.drawingsLayer.batchDraw();
+    };
+
+    const stopDrawing = () => {
+      if (!this.isEnabled) return;
+      if (!isDrawing || !currentLine) return;
+
+      // Store current line before resetting state
+      const lineToConvert = currentLine;
+
+      // Reset drawing state immediately so we can start a new stroke
+      isDrawing = false;
+      currentLine = null;
+
+      // Create a temporary layer to render the line
+      const tempLayer = new Konva.Layer({
+        clearBeforeDraw: true,
+      });
+      this.mediaLayerRef.getStage()?.add(tempLayer);
+      tempLayer.add(lineToConvert);
+
+      // Get the line canvas with transparent background
+      // Get the bounding box of the line
+      const lineBounds = lineToConvert.getClientRect();
+
+      // Create canvas with just enough size to contain the line
+      const lineCanvas = tempLayer.toCanvas({
+        x: lineBounds.x,
+        y: lineBounds.y,
+        width: lineBounds.width,
+        height: lineBounds.height,
+        pixelRatio: 1,
+      });
+
+      if (this.onDrawCallback) {
+        this.onDrawCallback(lineCanvas, lineBounds);
+      }
+
+      // Clean up
+      lineToConvert.destroy();
+      tempLayer.destroy();
+
+      this.drawingsLayer.batchDraw();
+    };
+
+    // Check if point is within capture canvas bounds
+    const isWithinCaptureCanvas = (pos: { x: number; y: number }) => {
+      const captureBox = this.captureCanvas.getClientRect();
+      return (
+        pos.x >= captureBox.x &&
+        pos.x <= captureBox.x + captureBox.width &&
+        pos.y >= captureBox.y &&
+        pos.y <= captureBox.y + captureBox.height
+      );
+    };
+
+    // Add event listeners
+    const stage = this.mediaLayerRef.getStage();
+    if (!stage) return;
+
+    stage.on("mousedown touchstart", (e) => {
+      const pos = stage.getPointerPosition();
+      if (pos && isWithinCaptureCanvas(pos)) {
+        startDrawing(pos);
+      }
+    });
+
+    stage.on("mousemove touchmove", (e) => {
+      const pos = stage.getPointerPosition();
+      if (pos && isWithinCaptureCanvas(pos)) {
+        draw(pos);
+      }
+    });
+
+    stage.on("mouseup touchend", async () => {
+      stopDrawing();
+      await this.render();
+    });
+
+    // Store cleanup function
+    this.cleanupFunction = () => {
+      stage.off("mousedown touchstart");
+      stage.off("mousemove touchmove");
+      stage.off("mouseup touchend");
+    };
+  }
+
+  public enablePaintMode() {
+    this.isEnabled = true;
+    if (!this.cleanupFunction) {
+      this.paintMode();
+    }
+  }
+
+  public enableDragging() {
+    // Enable dragging for all nodes in media layer
+    this.imageNodes?.forEach((node) => {
+      node.kNode.draggable(true);
+      node.kNode.listening(true);
+    });
+    this.mediaLayerRef.batchDraw();
+  }
+
+  public disableDragging() {
+    // Disable dragging for all nodes in media layer
+    this.imageNodes?.forEach((node) => {
+      node.kNode.draggable(false);
+      node.kNode.listening(false);
+    });
+    this.mediaLayerRef.batchDraw();
+  }
+
+  public disablePaintMode() {
+    this.isEnabled = false;
+    if (this.cleanupFunction) {
+      this.cleanupFunction();
+      this.cleanupFunction = null;
+    }
   }
 
   public previewCopyListener() {
-    this.previewCanvas.on('mousedown touchstart', () => {
+    this.previewCanvas.on("mousedown touchstart", () => {
       if (!this.outputBitmap) {
-        console.log("No preview image to copy"); 
+        console.log("No preview image to copy");
         return;
       }
-      
-      // TODO create it as a node.
-      
-      // const imageNode = new ImageNode({
-      //   mediaLayerRef: this.mediaLayer,
-      //   canvasPosition: this.renderEngine.captureCanvas.position(),
-      //   canvasSize: this.renderEngine.captureCanvas.size(),
-      //   imageFile: imageFile,
-      //   selectionManagerRef: this.selectionManager,
-      // });
 
       // Create draggable preview copy
       const previewCopy = new Konva.Image({
@@ -148,43 +439,49 @@ export class RealTimeDrawEngine {
         height: this.height,
         image: this.outputBitmap,
         draggable: true,
-        listening: true
+        listening: true,
       });
 
-      this.mediaLayerRef.add(previewCopy);
-      previewCopy.moveToTop();
-      this.mediaLayerRef.draw();
+      // this.drawingsLayer.add(previewCopy); // Add to drawingsLayer instead of mediaLayerRef
+      // this.drawingsLayer.batchDraw();
+
+      // Invoke the callback with the preview copy
+      if (this.onPreviewCopyCallback) {
+        this.onPreviewCopyCallback(previewCopy);
+      }
 
       // Start dragging immediately
       previewCopy.startDrag();
 
       // Handle drag events
-      previewCopy.on('dragmove', () => {
-        this.mediaLayerRef.draw();
+      previewCopy.on("dragmove", () => {
+        // this.drawingsLayer.draw();
       });
 
-      previewCopy.on('dragend', () => {
+      previewCopy.on("dragend", () => {
         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()
+            y: this.captureCanvas.y(),
           });
-          this.mediaLayerRef.batchDraw();
+          //this.drawingsLayer.batchDraw();
         } else {
           // Remove if dropped outside capture area
-          previewCopy.destroy();
-          this.mediaLayerRef.batchDraw();
+          //previewCopy.destroy();
+          //this.drawingsLayer.batchDraw();
         }
       });
     });
   }
 
-  public findImageNodeById(id: string): (ImageNode | TextNode | ShapeNode | undefined) {
-    return this.imageNodes.find(node => {
+  public findImageNodeById(
+    id: string,
+  ): ImageNode | TextNode | ShapeNode | PaintNode | undefined {
+    return this.imageNodes.find((node) => {
       if (node.kNode) {
         return node.kNode.id() === id;
       }
@@ -327,23 +624,33 @@ export class RealTimeDrawEngine {
   private handleNodeDragEnd = async () => {
     // Clean up any existing state
     if (this.isProcessing) {
-      console.log("isProcessing Returning")
+      console.log("isProcessing Returning");
       return;
     }
     console.log("Node drag ended");
     this.isProcessing = true;
     await this.render();
   };
-  
+
   public async addNodes(node: MediaNode) {
-   
-    if (node instanceof ImageNode || node instanceof TextNode || node instanceof ShapeNode) {
+    if (
+      node instanceof ImageNode ||
+      node instanceof TextNode ||
+      node instanceof ShapeNode ||
+      node instanceof PaintNode
+    ) {
       console.debug("Adding node:", node);
       this.imageNodes.push(node);
-      console.log(this.imageNodes)
-      node.kNode.on("dragend", this.handleNodeDragEnd);
+      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();
+    }
+
+    await this.render();
   }
 
   public removeNodes(node: MediaNode) {
@@ -353,7 +660,12 @@ export class RealTimeDrawEngine {
         node.kNode.off("dragend", this.handleNodeDragEnd);
         this.videoNodes.splice(index, 1);
       }
-    } else if (node instanceof ImageNode || node instanceof TextNode) {
+    } else if (
+      node instanceof ImageNode ||
+      node instanceof TextNode ||
+      node instanceof ShapeNode ||
+      node instanceof PaintNode
+    ) {
       const index = this.imageNodes.indexOf(node);
       if (index > -1) {
         node.kNode.off("dragend", this.handleNodeDragEnd);
@@ -377,7 +689,6 @@ export class RealTimeDrawEngine {
     try {
       const box = config.layerOfInterest.getClientRect();
       const stage = config.layerOfInterest.getStage();
-      
 
       const x = config.x !== undefined ? config.x : Math.floor(box.x);
       const y = config.y !== undefined ? config.y : Math.floor(box.y);
@@ -463,7 +774,7 @@ export class RealTimeDrawEngine {
     base64String: string,
   ): Promise<ImageBitmap> {
     // Create an image element
-    const img = new Image();
+    const img = document.createElement("img");
 
     // Convert base64 to data URL if it doesn't include the prefix
     const dataUrl = base64String.startsWith("data:")
@@ -514,7 +825,7 @@ export class RealTimeDrawEngine {
 
   public async render() {
     // only pick nodes that intersect wi th the canvas on screen bounds to freeze.
-    
+
     this.mediaLayerRef.draw();
     // Output all nodes in mediaLayerRef
     const nodes = this.mediaLayerRef.getChildren();

frontend/realtime/src/KonvaApp/UndoRedo/CreateCommand.ts

@@ -6,7 +6,7 @@ import {
   NodeTransformer,
   SelectionManager,
 } from "../NodesManagers";
-import { RenderEngine } from "../RenderingPrimitives/RenderEngine";
+import { RealTimeDrawEngine } from "../RenderingPrimitives/RealTimeDrawEngine";
 
 import { toolbarNode } from "~/signals/uiAccess/toolbarNode";
 
@@ -19,7 +19,7 @@ export class CreateCommand implements ICommand {
   private nodeTransformerRef: NodeTransformer;
   private selectionManagerRef: SelectionManager;
 
-  private renderEngineRef: RenderEngine;
+  private renderEngineRef: RealTimeDrawEngine;
 
   constructor({
     nodes,
@@ -33,7 +33,7 @@ export class CreateCommand implements ICommand {
     mediaLayerRef: Konva.Layer;
     nodesManagerRef: NodesManager;
     nodeTransformerRef: NodeTransformer;
-    renderEngineRef: RenderEngine;
+    renderEngineRef: RealTimeDrawEngine;
     selectionManagerRef: SelectionManager;
   }) {
     this.nodes = nodes;
@@ -53,6 +53,7 @@ export class CreateCommand implements ICommand {
       this.renderEngineRef.addNodes(node);
     });
     this.mediaLayerRef.draw();
+
   }
 
   undo() {
@@ -71,5 +72,7 @@ export class CreateCommand implements ICommand {
     this.nodeTransformerRef.clear();
     toolbarNode.hide();
     this.mediaLayerRef.draw();
+    this.renderEngineRef.render();
+ 
   }
 }

frontend/realtime/src/KonvaApp/UndoRedo/UndoRedoManager.ts

@@ -5,6 +5,11 @@ export class UndoStackManager {
   private undoStack: ICommand[] = [];
   private redoStack: ICommand[] = [];
   private readonly STACK_LIMIT = 20;
+  private onCommandExecuted: () => void;
+
+  constructor(onCommandExecuted: () => void) {
+    this.onCommandExecuted = onCommandExecuted;
+  }
 
   setDisabled(newState: boolean) {
     this.disabled = newState;
@@ -50,6 +55,7 @@ export class UndoStackManager {
       }
       console.log("undo:", this.undoStack);
     }
+    this.onCommandExecuted();
   }
 
   redo() {
@@ -68,5 +74,6 @@ export class UndoStackManager {
 
       console.log("Redo:", this.redoStack);
     }
+    this.onCommandExecuted();
   }
 }

frontend/realtime/src/KonvaApp/types/NodeData.ts

@@ -3,10 +3,11 @@ import { NodeType, BaseNode, ImageNode, TextNode, VideoNode, ShapeNode } from ".
 import { TransformationData } from "./Transformation";
 
 import { TextNodeData } from "./Text";
+import { PaintNode } from "../Nodes/PaintNode";
 
 export type MediaNode =
   // | NetworkedNode
-  BaseNode | VideoNode | ImageNode | TextNode | ShapeNode;
+  BaseNode | VideoNode | ImageNode | TextNode | ShapeNode | PaintNode;
 
 export type ImageNodeData = {
   mediaFileUrl: string;

frontend/realtime/src/PyServer/ServerClient.ts

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

frontend/realtime/src/PyServer/ServerTypes.ts

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

frontend/realtime/src/components/ImageGenerator.tsx

@@ -0,0 +1,524 @@
+import React, { useEffect, useState, useRef } from "react";
+import {
+  ServerClient,
+  fileToBase64,
+  base64ToImageUrl,
+} from "../PyServer/ServerClient";
+import {
+  ProgressUpdateMessage,
+  SetupResponse,
+  GenerateResponse,
+  StatusResponse,
+  UpdateSettingsResponse,
+} from "../PyServer/ServerTypes";
+
+// Default server URL
+const DEFAULT_SERVER_URL = "ws://localhost:8765";
+
+export const ImageGenerator: React.FC = () => {
+  // Server connection state
+  const [isConnected, setIsConnected] = useState(false);
+  const [serverUrl, setServerUrl] = useState(DEFAULT_SERVER_URL);
+  const clientRef = useRef<ServerClient | null>(null);
+
+  // Model state
+  const [modelPath, setModelPath] = useState("");
+  const [loraPath, setLoraPath] = useState("");
+  const [isModelInitialized, setIsModelInitialized] = useState(false);
+
+  // Generation parameters
+  const [sourceImage, setSourceImage] = useState<File | null>(null);
+  const [sourceImageUrl, setSourceImageUrl] = useState<string | null>(null);
+  const [prompt, setPrompt] = useState("");
+  const [loraStrength, setLoraStrength] = useState(1.0);
+  const [imageToImageStrength, setImageToImageStrength] = useState(0.75);
+  const [steps, setSteps] = useState(4);
+  const [guidanceScale, setGuidanceScale] = useState(1.0);
+  const [width, setWidth] = useState(1024);
+  const [height, setHeight] = useState(1024);
+
+  // Results and UI state
+  const [generatedImageUrl, setGeneratedImageUrl] = useState<string | null>(
+    null,
+  );
+  const [isLoading, setIsLoading] = useState(false);
+  const [progress, setProgress] = useState<ProgressUpdateMessage | null>(null);
+  const [errorMessage, setErrorMessage] = useState<string | null>(null);
+
+  // Initialize the client
+  useEffect(() => {
+    const client = new ServerClient(serverUrl, {
+      onOpen: () => {
+        console.log("Connected to server");
+        setIsConnected(true);
+        setErrorMessage(null);
+
+        // Request status to see if model is already initialized
+        client.requestStatus().catch((error) => {
+          console.error("Error requesting status:", error);
+        });
+      },
+
+      onClose: () => {
+        console.log("Disconnected from server");
+        setIsConnected(false);
+        setIsModelInitialized(false);
+      },
+
+      onError: (error) => {
+        console.error("Error:", error);
+        setErrorMessage(error.error);
+        setIsLoading(false);
+      },
+
+      onProgressUpdate: (update) => {
+        console.log("Progress:", update);
+        setProgress(update);
+
+        if (update.stage !== "error") {
+          setErrorMessage(null);
+        }
+      },
+
+      onSetupResponse: (response) => {
+        console.log("Setup response:", response);
+        setIsLoading(false);
+
+        if (response.success) {
+          setIsModelInitialized(true);
+          setErrorMessage(null);
+        } else {
+          setErrorMessage(response.error || "Unknown error during setup");
+        }
+      },
+
+      onUpdateSettingsResponse: (response) => {
+        console.log("Update settings response:", response);
+        setIsLoading(false);
+
+        if (response.success) {
+          setErrorMessage(null);
+        } else {
+          setErrorMessage(response.error || "Unknown error updating settings");
+        }
+      },
+
+      onGenerateResponse: (response) => {
+        console.log("Generate response:", response);
+        setIsLoading(false);
+
+        if (response.success && response.image) {
+          setGeneratedImageUrl(base64ToImageUrl(response.image));
+          setErrorMessage(null);
+        } else {
+          setErrorMessage(response.error || "Unknown error generating image");
+        }
+      },
+
+      onStatusResponse: (response) => {
+        console.log("Status response:", response);
+
+        if (response.success) {
+          setIsModelInitialized(response.model_initialized);
+
+          if (response.current_settings) {
+            // Update UI with current settings
+            setModelPath(response.current_settings.sdxl_checkpoint_path || "");
+            setLoraPath(response.current_settings.lora_path || "");
+
+            // Update generation parameters if they exist
+            if (response.current_settings.default_lora_strength !== undefined) {
+              setLoraStrength(response.current_settings.default_lora_strength);
+            }
+
+            if (
+              response.current_settings.default_image_to_image_strength !==
+              undefined
+            ) {
+              setImageToImageStrength(
+                response.current_settings.default_image_to_image_strength,
+              );
+            }
+
+            if (response.current_settings.default_steps !== undefined) {
+              setSteps(response.current_settings.default_steps);
+            }
+
+            if (
+              response.current_settings.default_guidance_scale !== undefined
+            ) {
+              setGuidanceScale(
+                response.current_settings.default_guidance_scale,
+              );
+            }
+
+            if (response.current_settings.default_width !== undefined) {
+              setWidth(response.current_settings.default_width);
+            }
+
+            if (response.current_settings.default_height !== undefined) {
+              setHeight(response.current_settings.default_height);
+            }
+          }
+        }
+      },
+    });
+
+    clientRef.current = client;
+
+    // Connect to the server
+    client.connect().catch((error) => {
+      console.error("Error connecting to server:", error);
+      setErrorMessage(`Failed to connect to server: ${error.message}`);
+    });
+
+    return () => {
+      client.disconnect();
+      clientRef.current = null;
+    };
+  }, [serverUrl]);
+
+  // Handle source image selection
+  const handleImageChange = async (e: React.ChangeEvent<HTMLInputElement>) => {
+    const file = e.target.files?.[0];
+    if (file) {
+      setSourceImage(file);
+      setSourceImageUrl(URL.createObjectURL(file));
+    }
+  };
+
+  // Handle setup submission
+  const handleSetup = async () => {
+    if (!clientRef.current) return;
+
+    setIsLoading(true);
+    setErrorMessage(null);
+
+    try {
+      await clientRef.current.setup(modelPath, loraPath || undefined);
+    } catch (error) {
+      console.error("Error setting up model:", error);
+      setErrorMessage(`Error setting up model: ${(error as Error).message}`);
+      setIsLoading(false);
+    }
+  };
+
+  // Handle update settings
+  const handleUpdateSettings = async () => {
+    if (!clientRef.current) return;
+
+    setIsLoading(true);
+    setErrorMessage(null);
+
+    try {
+      await clientRef.current.updateSettings({
+        lora_path: loraPath || null,
+        default_lora_strength: loraStrength,
+        default_image_to_image_strength: imageToImageStrength,
+        default_steps: steps,
+        default_guidance_scale: guidanceScale,
+        default_width: width,
+        default_height: height,
+      });
+    } catch (error) {
+      console.error("Error updating settings:", error);
+      setErrorMessage(`Error updating settings: ${(error as Error).message}`);
+      setIsLoading(false);
+    }
+  };
+
+  // Handle generation
+  const handleGenerate = async () => {
+    if (!clientRef.current || !sourceImage) {
+      setErrorMessage("Please select a source image first");
+      return;
+    }
+
+    setIsLoading(true);
+    setErrorMessage(null);
+    setGeneratedImageUrl(null);
+
+    try {
+      const base64Image = await fileToBase64(sourceImage);
+
+      await clientRef.current.generate({
+        image: base64Image,
+        prompt,
+        lora_strength: loraStrength,
+        image_to_image_strength: imageToImageStrength,
+        num_inference_steps: steps,
+        guidance_scale: guidanceScale,
+        generated_image_width: width,
+        generated_image_height: height,
+      });
+    } catch (error) {
+      console.error("Error generating image:", error);
+      setErrorMessage(`Error generating image: ${(error as Error).message}`);
+      setIsLoading(false);
+    }
+  };
+
+  return (
+    <div className="mx-auto max-w-4xl p-6">
+      <h1 className="mb-6 text-2xl font-bold">SDXL Image Generator</h1>
+
+      {/* Connection and Setup */}
+      <div className="mb-6 rounded-lg bg-gray-100 p-4">
+        <h2 className="mb-2 text-xl font-semibold">Server Connection</h2>
+        <div className="mb-4 flex gap-2">
+          <input
+            type="text"
+            value={serverUrl}
+            onChange={(e) => setServerUrl(e.target.value)}
+            className="flex-grow rounded border p-2"
+            placeholder="WebSocket URL"
+          />
+          <button
+            onClick={() => clientRef.current?.connect()}
+            disabled={isConnected}
+            className="rounded bg-blue-500 px-4 py-2 text-white disabled:bg-gray-400"
+          >
+            Connect
+          </button>
+          <button
+            onClick={() => clientRef.current?.disconnect()}
+            disabled={!isConnected}
+            className="rounded bg-red-500 px-4 py-2 text-white disabled:bg-gray-400"
+          >
+            Disconnect
+          </button>
+        </div>
+
+        <div className="mb-2 flex items-center">
+          <span className="mr-2">Status:</span>
+          <span
+            className={`rounded px-2 py-1 text-white ${isConnected ? "bg-green-500" : "bg-red-500"}`}
+          >
+            {isConnected ? "Connected" : "Disconnected"}
+          </span>
+          <span className="ml-4 mr-2">Model:</span>
+          <span
+            className={`rounded px-2 py-1 text-white ${isModelInitialized ? "bg-green-500" : "bg-yellow-500"}`}
+          >
+            {isModelInitialized ? "Initialized" : "Not Initialized"}
+          </span>
+        </div>
+      </div>
+
+      {/* Model Setup */}
+      <div className="mb-6 rounded-lg bg-gray-100 p-4">
+        <h2 className="mb-2 text-xl font-semibold">Model Setup</h2>
+        <div className="space-y-4">
+          <div>
+            <label className="mb-1 block">SDXL Checkpoint Path:</label>
+            <input
+              type="text"
+              value={modelPath}
+              onChange={(e) => setModelPath(e.target.value)}
+              className="w-full rounded border p-2"
+              placeholder="/path/to/sdxl/model"
+            />
+          </div>
+          <div>
+            <label className="mb-1 block">LoRA Path (optional):</label>
+            <input
+              type="text"
+              value={loraPath}
+              onChange={(e) => setLoraPath(e.target.value)}
+              className="w-full rounded border p-2"
+              placeholder="/path/to/lora"
+            />
+          </div>
+          <div className="flex gap-2">
+            <button
+              onClick={handleSetup}
+              disabled={!isConnected || !modelPath || isLoading}
+              className="rounded bg-green-500 px-4 py-2 text-white disabled:bg-gray-400"
+            >
+              Initialize Model
+            </button>
+            <button
+              onClick={handleUpdateSettings}
+              disabled={!isConnected || !isModelInitialized || isLoading}
+              className="rounded bg-purple-500 px-4 py-2 text-white disabled:bg-gray-400"
+            >
+              Update Settings
+            </button>
+          </div>
+        </div>
+      </div>
+
+      {/* Generation Parameters */}
+      <div className="mb-6 rounded-lg bg-gray-100 p-4">
+        <h2 className="mb-2 text-xl font-semibold">Generation Parameters</h2>
+        <div className="grid grid-cols-1 gap-4 md:grid-cols-2">
+          <div>
+            <label className="mb-1 block">Prompt:</label>
+            <textarea
+              value={prompt}
+              onChange={(e) => setPrompt(e.target.value)}
+              className="h-20 w-full rounded border p-2"
+              placeholder="Enter your prompt..."
+            />
+          </div>
+
+          <div>
+            <label className="mb-1 block">Source Image:</label>
+            <input
+              type="file"
+              accept="image/*"
+              onChange={handleImageChange}
+              className="mb-2"
+            />
+            {sourceImageUrl && (
+              <img
+                src={sourceImageUrl}
+                alt="Source"
+                className="max-h-40 rounded border"
+              />
+            )}
+          </div>
+
+          <div>
+            <label className="mb-1 block">
+              LoRA Strength: {loraStrength.toFixed(2)}
+            </label>
+            <input
+              type="range"
+              min="0"
+              max="1"
+              step="0.05"
+              value={loraStrength}
+              onChange={(e) => setLoraStrength(parseFloat(e.target.value))}
+              className="w-full"
+            />
+          </div>
+
+          <div>
+            <label className="mb-1 block">
+              Image-to-Image Strength: {imageToImageStrength.toFixed(2)}
+            </label>
+            <input
+              type="range"
+              min="0"
+              max="1"
+              step="0.05"
+              value={imageToImageStrength}
+              onChange={(e) =>
+                setImageToImageStrength(parseFloat(e.target.value))
+              }
+              className="w-full"
+            />
+          </div>
+
+          <div>
+            <label className="mb-1 block">Steps: {steps}</label>
+            <input
+              type="range"
+              min="1"
+              max="20"
+              step="1"
+              value={steps}
+              onChange={(e) => setSteps(parseInt(e.target.value))}
+              className="w-full"
+            />
+          </div>
+
+          <div>
+            <label className="mb-1 block">
+              Guidance Scale: {guidanceScale.toFixed(2)}
+            </label>
+            <input
+              type="range"
+              min="0"
+              max="10"
+              step="0.1"
+              value={guidanceScale}
+              onChange={(e) => setGuidanceScale(parseFloat(e.target.value))}
+              className="w-full"
+            />
+          </div>
+
+          <div>
+            <label className="mb-1 block">Width:</label>
+            <select
+              value={width}
+              onChange={(e) => setWidth(parseInt(e.target.value))}
+              className="w-full rounded border p-2"
+            >
+              <option value="512">512</option>
+              <option value="768">768</option>
+              <option value="1024">1024</option>
+              <option value="1280">1280</option>
+            </select>
+          </div>
+
+          <div>
+            <label className="mb-1 block">Height:</label>
+            <select
+              value={height}
+              onChange={(e) => setHeight(parseInt(e.target.value))}
+              className="w-full rounded border p-2"
+            >
+              <option value="512">512</option>
+              <option value="768">768</option>
+              <option value="1024">1024</option>
+              <option value="1280">1280</option>
+            </select>
+          </div>
+        </div>
+
+        <button
+          onClick={handleGenerate}
+          disabled={
+            !isConnected || !isModelInitialized || !sourceImage || isLoading
+          }
+          className="mt-4 rounded bg-blue-500 px-6 py-2 text-white disabled:bg-gray-400"
+        >
+          Generate Image
+        </button>
+      </div>
+
+      {/* Progress and Results */}
+      {isLoading && progress && (
+        <div className="mb-6 rounded-lg bg-gray-100 p-4">
+          <h2 className="mb-2 text-xl font-semibold">Progress</h2>
+          <div className="mb-2">
+            <div className="mb-1 flex justify-between">
+              <span>{progress.stage}</span>
+              <span>{progress.progress}%</span>
+            </div>
+            <div className="h-2.5 w-full rounded-full bg-gray-300">
+              <div
+                className="h-2.5 rounded-full bg-blue-600"
+                style={{ width: `${progress.progress}%` }}
+              ></div>
+            </div>
+            {progress.file && (
+              <div className="mt-1 text-sm text-gray-600">{progress.file}</div>
+            )}
+          </div>
+        </div>
+      )}
+
+      {errorMessage && (
+        <div className="mb-6 rounded border border-red-400 bg-red-100 px-4 py-3 text-red-700">
+          <p>{errorMessage}</p>
+        </div>
+      )}
+
+      {generatedImageUrl && (
+        <div className="mb-6 rounded-lg bg-gray-100 p-4">
+          <h2 className="mb-2 text-xl font-semibold">Generated Image</h2>
+          <div className="flex justify-center">
+            <img
+              src={generatedImageUrl}
+              alt="Generated"
+              className="max-w-full rounded border"
+            />
+          </div>
+        </div>
+      )}
+    </div>
+  );
+};

frontend/realtime/src/components/features/ToolbarMain/ToolbarMain.tsx

@@ -25,6 +25,8 @@ import { ToolbarMainButtonNames } from "./enum";
 import { dispatchUiEvents, uiEvents } from "~/signals/uiEvents";
 import { paintColor } from "~/signals/uiEvents/toolbarMain/paintMode";
 import { ColorPicker } from "~/components/ui/TextEditor/ColorPicker";
+import { PaintModeMenu } from "~/components/ui/ToolbarMain/PaintModeMenu";
+import { EraseModeMenu } from "~/components/ui/ToolbarMain/EraseModeMenu";
 
 export const ToolbarMain = ({
   disabled = false,
@@ -47,7 +49,7 @@ export const ToolbarMain = ({
           paperWrapperStyles,
           "glass",
           disabled &&
-            "pointer-events-none cursor-default bg-ui-border shadow-md",
+          "pointer-events-none cursor-default bg-ui-border shadow-md",
         )}
       >
         <div className="flex flex-col items-center gap-2">
@@ -155,9 +157,10 @@ export const ToolbarMain = ({
 
         <div className="flex flex-col items-center gap-2">
           <div className="flex flex-col items-center gap-2">
+            { /* Conditionally show paint menu on mode selection */}
             {buttonProps.PAINT.active ? (
               <div className="relative">
-                <ColorPicker
+                <PaintModeMenu
                   color={paintColor.value}
                   onChange={dispatchUiEvents.toolbarMain.setPaintColor}
                   faIcon={faPaintbrush}
@@ -182,11 +185,25 @@ export const ToolbarMain = ({
                 tooltip="Paint Brush"
               />
             )}
-            <ToolbarButton
-              icon={faEraser}
-              buttonProps={buttonProps.ERASER}
-              tooltip="Eraser"
-            />
+
+            { /* Conditionally show eraser menu on mode selection */}
+            {buttonProps.ERASER.active ? (
+              <EraseModeMenu
+                faIcon={faEraser}
+                borderStyle="border-2  bg-primary/30 border-2 border-primary hover:bg-primary/30 text-white"
+                defaultOpen={true}
+                anchor="right"
+                anchorGap={12}
+                closeOnMouseLeave={true}
+              />
+            )
+              : (
+                <ToolbarButton
+                  icon={faEraser}
+                  buttonProps={buttonProps.ERASER}
+                  tooltip="Eraser"
+                />
+              )}
           </div>
         </div>
 

frontend/realtime/src/components/features/ToolbarNode/ToolbarNode.tsx

@@ -13,9 +13,6 @@ import {
 import { paperWrapperStyles, toolTipStyles } from "~/components/styles";
 import { ToolbarNodeButtonNames } from "./enums";
 import { ToolbarNodeButtonData } from "./data";
-import { Popover, PopoverButton, PopoverPanel } from "@headlessui/react";
-import { Fragment } from "react/jsx-runtime";
-import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import { ColorPicker } from "~/components/ui/TextEditor/ColorPicker";
 
 export interface ToolbarNodeProps {

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

@@ -1,6 +1,8 @@
 import { Transition } from "@headlessui/react";
 import React, { useState, useRef, useEffect, ReactNode } from "react";
 import { twMerge } from "tailwind-merge";
+import { Button } from "./Button";
+import { faMinus, faPlus } from "@fortawesome/pro-solid-svg-icons";
 
 interface SliderProps {
   min: number;
@@ -13,6 +15,8 @@ interface SliderProps {
   tooltipContent?: ReactNode;
   suffix?: string;
   innerLabel?: string;
+  showDecrement?: boolean,
+  showIncrement?: boolean,
 }
 
 export const Slider = ({
@@ -26,6 +30,8 @@ export const Slider = ({
   tooltipContent,
   suffix,
   innerLabel,
+  showDecrement,
+  showIncrement
 }: SliderProps) => {
   const [localValue, setLocalValue] = useState(value);
   const sliderRef = useRef<HTMLDivElement>(null);
@@ -36,15 +42,19 @@ export const Slider = ({
     setLocalValue(value);
   }, [value]);
 
+  const setClampedValue = (value: number) => {
+    const roundedValue = Math.round(value / step) * step;
+    const clampedValue = Math.min(Math.max(roundedValue, min), max);
+    setLocalValue(clampedValue);
+    onChange(clampedValue);
+  }
+
   const updatePosition = (x: number) => {
     if (!sliderRef.current) return;
     const rect = sliderRef.current.getBoundingClientRect();
     const fraction = (x - rect.left) / rect.width;
     const newValue = fraction * (max - min) + min;
-    const roundedValue = Math.round(newValue / step) * step;
-    const clampedValue = Math.min(Math.max(roundedValue, min), max);
-    setLocalValue(clampedValue);
-    onChange(clampedValue);
+    setClampedValue(newValue);
   };
 
   const handleMouseDown = (e: React.MouseEvent<HTMLDivElement>) => {
@@ -81,49 +91,75 @@ export const Slider = ({
     }
   };
 
+  const handleDecrement = () => {
+    setClampedValue(localValue - step);
+  }
+
+  const handleIncrement = () => {
+    setClampedValue(localValue + step);
+  }
+
   const percentage = ((localValue - min) / (max - min)) * 100;
   const displayValue =
     localValue.toFixed(step.toString().split(".")[1]?.length || 0) +
     (suffix || "");
 
   return (
-    <div
-      ref={sliderRef}
-      className={twMerge(
-        "glass group relative h-7 w-full cursor-pointer overflow-hidden rounded-lg border border-ui-border",
-        isDragging && "!bg-ui-controls/90",
-        className,
+    <div className="flex w-full">
+      {showDecrement && (
+        <Button
+          icon={faMinus}
+          className="size-6 text-white/80 bg-transparent rounded-full hover:bg-white/10 focus-visible:outline-primary active:bg-primary/30 my-auto mr-1"
+          onClick={handleDecrement}
+        />
       )}
-      onMouseDown={handleMouseDown}
-    >
+
       <div
+        ref={sliderRef}
         className={twMerge(
-          "absolute h-7 bg-primary/30 transition-colors duration-300 group-hover:bg-primary/50",
-          isDragging && "!bg-primary/50",
+          "glass group relative h-7 w-full cursor-pointer overflow-hidden rounded-lg border border-ui-border",
+          isDragging && "!bg-ui-controls/90",
+          className,
         )}
-        style={{ width: `${percentage}%` }}
+        onMouseDown={handleMouseDown}
       >
-        {innerLabel && (
-          <span
-            className={twMerge(
-              "absolute top-1/2 ml-2.5 -translate-y-1/2 text-nowrap text-sm font-medium text-white/60 transition-colors duration-300 group-hover:text-white",
-              isDragging && "!text-white",
-            )}
-          >
-            {innerLabel}
-          </span>
-        )}
         <div
           className={twMerge(
-            "absolute right-0 top-1/2 mr-1.5 h-3.5 w-0.5 -translate-y-1/2",
-            isDragging ? "bg-white" : "bg-white/50",
+            "absolute h-7 bg-primary/30 transition-colors duration-300 group-hover:bg-primary/50",
+            isDragging && "!bg-primary/50",
           )}
-          onMouseDown={handleMouseDown}
-          onMouseEnter={handleMouseEnter}
-          onMouseLeave={handleMouseLeave}
-        />
+          style={{ width: `${percentage}%` }}
+        >
+          {innerLabel && (
+            <span
+              className={twMerge(
+                "absolute top-1/2 ml-2.5 -translate-y-1/2 text-nowrap text-sm font-medium text-white/60 transition-colors duration-300 group-hover:text-white",
+                isDragging && "!text-white",
+              )}
+            >
+              {innerLabel}
+            </span>
+          )}
+          <div
+            className={twMerge(
+              "absolute right-0 top-1/2 mr-1.5 h-3.5 w-0.5 -translate-y-1/2",
+              isDragging ? "bg-white" : "bg-white/50",
+            )}
+            onMouseDown={handleMouseDown}
+            onMouseEnter={handleMouseEnter}
+            onMouseLeave={handleMouseLeave}
+          />
+        </div>
       </div>
 
+      {showIncrement && (
+        <Button
+          icon={faPlus}
+          className="size-6 text-white/80 bg-transparent rounded-full hover:bg-white/10 focus-visible:outline-primary active:bg-primary/30 my-auto ml-1"
+          onClick={handleIncrement}
+        />
+      )}
+
       {showTooltip && (
         <Transition
           as={"div"}

frontend/realtime/src/components/ui/TextEditor/ColorPicker.tsx

@@ -1,5 +1,3 @@
-import { ChangeEvent, useState, useEffect } from "react";
-import { twMerge } from "tailwind-merge";
 import {
   faFont,
   faHashtag,
@@ -7,11 +5,13 @@ import {
 } from "@fortawesome/pro-solid-svg-icons";
 import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import { Popover, PopoverButton, PopoverPanel } from "@headlessui/react";
-import { HexColorPicker } from "react-colorful";
+import { ChangeEvent, useEffect, useState } from "react";
+import { HexAlphaColorPicker } from "react-colorful";
+import { twMerge } from "tailwind-merge";
 
+import { paperWrapperStyles } from "~/components/styles";
 import { Button } from "../Button";
 import { Input } from "../Input";
-import { paperWrapperStyles } from "~/components/styles";
 
 export const ColorPicker = ({
   color: prevColor,
@@ -119,7 +119,7 @@ export const ColorPicker = ({
                 }
                 static
               >
-                <HexColorPicker
+                <HexAlphaColorPicker
                   className="overflow-hidden"
                   color={currColor}
                   onChange={(color) => {

frontend/realtime/src/components/ui/ToolbarMain/EraseModeMenu.tsx

@@ -0,0 +1,102 @@
+import {
+  faFont,
+  faHashtag,
+  IconDefinition,
+} from "@fortawesome/pro-solid-svg-icons";
+import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
+import { Popover, PopoverButton, PopoverPanel } from "@headlessui/react";
+import { ChangeEvent, useEffect, useState } from "react";
+import { HexColorPicker } from "react-colorful";
+import { twMerge } from "tailwind-merge";
+
+import { useSignals } from "@preact/signals-react/runtime";
+import { paperWrapperStyles } from "~/components/styles";
+import { BRUSH_MAX_SIZE, BRUSH_MIN_SIZE, paintBrushSize, setPaintBrushSize } from "~/signals/uiEvents/toolbarMain/paintMode";
+import { Button } from "../Button";
+import { Input } from "../Input";
+import { Slider } from "../Slider";
+import { eraseBrushSize, setEraseBrushSize } from "~/signals/uiEvents/toolbarMain/eraseMode";
+
+export const EraseModeMenu = ({
+  faIcon = faFont,
+  borderStyle,
+  anchor = "bottom",
+  anchorGap = 0,
+  defaultOpen = false,
+  closeOnMouseLeave = false,
+}: {
+  faIcon?: IconDefinition;
+  borderStyle?: string;
+  anchor?: "top" | "right" | "bottom" | "left";
+  anchorGap?: number;
+  defaultOpen?: boolean;
+  closeOnMouseLeave?: boolean;
+}) => {
+  const [isOpen, setIsOpen] = useState(defaultOpen);
+
+  useEffect(() => {
+    setIsOpen(defaultOpen);
+  }, [defaultOpen]);
+
+  const handleMouseLeave = (close: () => void) => {
+    if (closeOnMouseLeave) {
+      setIsOpen(false);
+      close();
+    }
+  };
+
+  // Rerender the component when these signals change:
+  useSignals();
+  const brushSize = eraseBrushSize.value;
+
+  return (
+    <Popover className="relative">
+      {({ open, close }) => (
+        <>
+          <PopoverButton
+            className={twMerge(
+              "flex size-10 items-center justify-center gap-1 rounded-lg bg-ui-controls p-2",
+              borderStyle,
+            )}
+            onMouseEnter={() => defaultOpen && setIsOpen(true)}
+          >
+            <FontAwesomeIcon
+              icon={faIcon}
+            />
+          </PopoverButton>
+          {(open || isOpen) && (
+            <div onMouseLeave={() => handleMouseLeave(close)}>
+              <PopoverPanel
+                anchor={anchor}
+                className={twMerge(
+                  paperWrapperStyles,
+                  "flex flex-col items-center gap-2 overflow-hidden w-[198px]",
+                )}
+                style={
+                  {
+                    "--anchor-gap": `${anchorGap}px`,
+                  } as React.CSSProperties
+                }
+                static
+              >
+                <div className="flex flex-col w-full gap-2">
+                  <p className="text-white font-medium text-sm w-full justify-start">Eraser Size:</p>
+                  <Slider
+                    min={BRUSH_MIN_SIZE}
+                    max={BRUSH_MAX_SIZE}
+                    value={brushSize}
+                    onChange={setEraseBrushSize}
+                    step={1}
+                    innerLabel={brushSize.toString()}
+                    showDecrement
+                    showIncrement
+                  />
+                </div>
+              </PopoverPanel>
+            </div>
+          )}
+        </>
+      )}
+    </Popover>
+  );
+};

frontend/realtime/src/components/ui/ToolbarMain/PaintModeMenu.tsx

@@ -0,0 +1,194 @@
+import {
+  faFont,
+  faHashtag,
+  IconDefinition,
+} from "@fortawesome/pro-solid-svg-icons";
+import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
+import { Popover, PopoverButton, PopoverPanel } from "@headlessui/react";
+import { ChangeEvent, useEffect, useState } from "react";
+import { HexAlphaColorPicker } from "react-colorful";
+import { twMerge } from "tailwind-merge";
+
+import { useSignals } from "@preact/signals-react/runtime";
+import { paperWrapperStyles } from "~/components/styles";
+import { BRUSH_MAX_SIZE, BRUSH_MIN_SIZE, paintBrushSize, setPaintBrushSize } from "~/signals/uiEvents/toolbarMain/paintMode";
+import { Button } from "../Button";
+import { Input } from "../Input";
+import { Slider } from "../Slider";
+
+export const PaintModeMenu = ({
+  color: prevColor,
+  onChange,
+  faIcon = faFont,
+  borderStyle,
+  showBar = true,
+  anchor = "bottom",
+  anchorGap = 0,
+  fillBg = false,
+  staticIconColor,
+  streamChanges = false,
+  defaultOpen = false,
+  closeOnMouseLeave = false,
+}: {
+  color: string;
+  onChange: (newColor: string) => void;
+  faIcon?: IconDefinition;
+  borderStyle?: string;
+  showBar?: boolean;
+  anchor?: "top" | "right" | "bottom" | "left";
+  anchorGap?: number;
+  fillBg?: boolean;
+  staticIconColor?: string;
+  streamChanges?: boolean;
+  defaultOpen?: boolean;
+  closeOnMouseLeave?: boolean;
+}) => {
+  const [isOpen, setIsOpen] = useState(defaultOpen);
+  const [{ currColor, textInput }, setStates] = useState<{
+    currColor: string;
+    textInput: string;
+  }>({
+    currColor: prevColor,
+    textInput: prevColor.substring(1),
+  });
+
+  useEffect(() => {
+    setIsOpen(defaultOpen);
+  }, [defaultOpen]);
+
+  const handleMouseLeave = (close: () => void) => {
+    if (closeOnMouseLeave) {
+      onChange(currColor);
+      setIsOpen(false);
+      close();
+    }
+  };
+
+  const handleTextInput = (e: ChangeEvent<HTMLInputElement>) => {
+    if (isHexColorCode(e.target.value)) {
+      setStates({
+        currColor: `#${e.target.value}`,
+        textInput: e.target.value,
+      });
+    } else {
+      setStates((curr) => ({
+        ...curr,
+        textInput: e.target.value,
+      }));
+    }
+  };
+
+  const isHexColorCode = (color: string) => {
+    const hexColorRegex = /^(?:[0-9a-fA-F]{3}){1,2}$/;
+    return hexColorRegex.test(color);
+  };
+
+  // Rerender the component when these signals change:
+  useSignals();
+  const brushSize = paintBrushSize.value;
+
+  return (
+    <Popover className="relative">
+      {({ open, close }) => (
+        <>
+          <PopoverButton
+            className={twMerge(
+              "flex size-10 items-center gap-1 rounded-lg bg-ui-controls p-2",
+              borderStyle,
+              showBar ? "flex-col" : "justify-center",
+            )}
+            style={fillBg ? { backgroundColor: prevColor } : {}}
+            onMouseEnter={() => defaultOpen && setIsOpen(true)}
+          >
+            <FontAwesomeIcon
+              icon={faIcon}
+              color={staticIconColor ?? (fillBg ? undefined : prevColor)}
+            />
+            {showBar && (
+              <span
+                className="h-1 w-full"
+                style={{ backgroundColor: prevColor }}
+              />
+            )}
+          </PopoverButton>
+          {(open || isOpen) && (
+            <div onMouseLeave={() => handleMouseLeave(close)}>
+              <PopoverPanel
+                anchor={anchor}
+                className={twMerge(
+                  paperWrapperStyles,
+                  "flex flex-col items-center gap-2 overflow-hidden",
+                )}
+                style={
+                  {
+                    "--anchor-gap": `${anchorGap}px`,
+                  } as React.CSSProperties
+                }
+                static
+              >
+                <HexAlphaColorPicker
+                  className="overflow-hidden"
+                  color={currColor}
+                  onChange={(color) => {
+                    setStates({
+                      currColor: color,
+                      textInput: color.substring(1),
+                    });
+                    if (streamChanges) {
+                      onChange(color);
+                    }
+                  }}
+                />
+                <Input
+                  style={{ width: "198px" }}
+                  icon={faHashtag}
+                  value={textInput}
+                  onChange={handleTextInput}
+                />
+                <div className="flex flex-col w-full mt-2 gap-2">
+                  <p className="text-white font-medium text-sm w-full justify-start">Brush Size:</p>
+                  <Slider
+                    min={BRUSH_MIN_SIZE}
+                    max={BRUSH_MAX_SIZE}
+                    value={brushSize}
+                    onChange={setPaintBrushSize}
+                    step={1}
+                    innerLabel={brushSize.toString()}
+                    showDecrement
+                    showIncrement
+                  />
+                </div>
+                {!streamChanges && (
+                  <div className="flex w-full justify-center gap-2">
+                    <Button
+                      variant="secondary"
+                      onClick={() => {
+                        setStates({
+                          currColor: prevColor,
+                          textInput: prevColor.substring(1),
+                        });
+                        setIsOpen(false);
+                        close();
+                      }}
+                    >
+                      Cancel
+                    </Button>
+                    <Button
+                      onClick={() => {
+                        onChange(currColor);
+                        setIsOpen(false);
+                        close();
+                      }}
+                    >
+                      OK
+                    </Button>
+                  </div>
+                )}
+              </PopoverPanel>
+            </div>
+          )}
+        </>
+      )}
+    </Popover>
+  );
+};

frontend/realtime/src/signals/uiEvents/toolbarMain/eraseMode.ts

@@ -0,0 +1,14 @@
+import { effect, signal } from "@preact/signals-react";
+import { clamp, normalize } from "~/utilities";
+import { BRUSH_MAX_SIZE, BRUSH_MIN_SIZE, DEFAULT_PAINT_BRUSH_SIZE, paintBrushSize } from "./paintMode";
+
+export const eraseBrushSize = signal<number>(DEFAULT_PAINT_BRUSH_SIZE);
+export const setEraseBrushSize = (brushSize: number) => {
+  eraseBrushSize.value = clamp(brushSize, BRUSH_MIN_SIZE, BRUSH_MAX_SIZE);
+}
+
+export const onEraseBrushSizeChanged = (callback: (normalizedSize: number) => void) => {
+  effect(() => {
+    callback(normalize(paintBrushSize.value, BRUSH_MIN_SIZE, BRUSH_MAX_SIZE));
+  });
+};

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

@@ -3,14 +3,18 @@ import {
   loadingBarRetryDispatch,
   loadingBarRetryEventHandler,
 } from "./loadingBar";
-import { onPaintColorChanged, setPaintColor } from "./paintMode";
+import { onEraseBrushSizeChanged, setEraseBrushSize } from "./eraseMode";
+import { onPaintColorChanged, setPaintBrushSize, setPaintColor } from "./paintMode";
 export const dispatchers = {
   ...buttonDispatchers,
   loadingBarRetry: loadingBarRetryDispatch,
   setPaintColor,
+  setPaintBrushSize,
+  setEraseBrushSize
 };
 export const eventsHandlers = {
   ...buttonEventsHandlers,
   loadingBarRetry: loadingBarRetryEventHandler,
-  onPaintColorChanged
+  onPaintColorChanged,
+  onEraseBrushSizeChanged
 };

frontend/realtime/src/signals/uiEvents/toolbarMain/paintMode.ts

@@ -1,15 +1,22 @@
-// Prompt settings
-// - Prompt text
-
 import { effect, signal } from "@preact/signals-react";
+import { clamp, normalize } from "~/utilities";
 
-// - Prompt strength
 export const DEFAULT_PAINT_COLOR: string = "#ff0000";
 export const paintColor = signal<string>(DEFAULT_PAINT_COLOR);
 export const setPaintColor = (data: string) => {
   paintColor.value = data;
 };
 
+export const BRUSH_MIN_SIZE: number = 0;
+export const BRUSH_MAX_SIZE: number = 100;
+// NOTE: Must be between MIN and MAX sizes - there's no runtime checks
+export const DEFAULT_PAINT_BRUSH_SIZE: number = 50;
+export const paintBrushSize = signal<number>(DEFAULT_PAINT_BRUSH_SIZE);
+export const setPaintBrushSize = (brushSize: number) => {
+  paintBrushSize.value = clamp(brushSize, BRUSH_MIN_SIZE, BRUSH_MAX_SIZE);
+  console.debug(`Paint brush size set to ${paintBrushSize.value}`);
+}
+
 export const onPaintColorChanged = (callback: (data: string) => void) => {
   effect(() => {
     if (paintColor.value) {
@@ -17,3 +24,9 @@ export const onPaintColorChanged = (callback: (data: string) => void) => {
     }
   });
 };
+
+export const onPaintBrushSizeChanged = (callback: (normalizedSize: number) => void) => {
+  effect(() => {
+    callback(normalize(paintBrushSize.value, BRUSH_MIN_SIZE, BRUSH_MAX_SIZE));
+  });
+};

frontend/realtime/src/types/ServerTypes.ts

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

frontend/realtime/src/utilities/index.ts

@@ -10,3 +10,9 @@ export const getFileName = (file: File) => {
 export const getFileExtension = (file: File) => {
   return file.name.substring(file.name.lastIndexOf(".") + 1);
 };
+export const clamp = (num: number, min: number, max: number) => {
+  return Math.min(Math.max(num, min), max);
+}
+export const normalize = (num: number, min: number, max: number) => {
+  return (num - min) / (max - min);
+}

frontend/realtime/start_server.ps1

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