Things flow nicely a bunch of bug fixes.

frontend/realtime/paint_py_server/paint_test.py

@@ -0,0 +1,211 @@
+import asyncio
+import base64
+import json
+import logging
+import pathlib
+import time
+from io import BytesIO
+from typing import Dict, Any, Optional, Tuple, Callable, List
+from dataclasses import dataclass
+from dataclasses_json import dataclass_json
+from enum import Enum, auto
+import base64
+from PIL import Image
+import io
+import torch
+from diffusers import AutoPipelineForImage2Image, LCMScheduler
+from diffusers.utils import make_image_grid, load_image
+
+import torch
+import websockets
+from PIL import Image
+from diffusers import (
+    StableDiffusionXLPipeline,
+    StableDiffusionXLImg2ImgPipeline,
+    AutoPipelineForImage2Image,
+    LCMScheduler,
+    UNet2DConditionModel,
+)
+import torch
+import torchvision.transforms as transforms
+
+transform = transforms.Compose([
+    transforms.ToTensor(),  # This converts PIL Image to tensor and scales to [0, 1]
+])
+
+@dataclass_json
+@dataclass
+class ProgressMessage:
+    message: str
+    progress: float
+    type: str = "progress"
+    error: Optional[str] = None
+
+@dataclass_json
+@dataclass
+class LoadModelRequest:
+    model_path: str
+    lora_path: Optional[str] = None
+    command: str = "load_model"
+
+@dataclass_json
+@dataclass
+class GenerateRequest:
+    image: str
+    prompt: str
+    command: str = "generate"
+    strength: float = 0.6
+    guidance_scale: float = 2.0
+    num_inference_steps: int = 4
+
+@dataclass_json
+@dataclass
+class GenerateResponse:
+    image: str
+    type: str = "result"
+
+def save_image(image: Image.Image, base_path: str = "outputs") -> str:
+    """Save a PIL Image to disk with a timestamp-based filename."""
+    try:
+        # Create outputs directory if it doesn't exist
+        output_dir = pathlib.Path(base_path)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Generate filename with timestamp
+        timestamp = time.strftime("%Y%m%d-%H%M%S")
+        filename = f"generated-{timestamp}.png"
+        filepath = output_dir / filename
+        
+        # Save the image
+        print(f"Saving image to: {filepath}")
+        image.save(str(filepath), format="PNG")
+        print(f"Successfully saved image to: {filepath}")
+        return str(filepath)
+    except Exception as e:
+        print(f"Error saving image: {e}")
+        return ""
+
+
+async def send_progress(websocket, message: str, progress: float = 0.0, error: str = None):
+    """Send progress updates to the client."""
+    progress_msg = ProgressMessage(
+        message=message,
+        progress=progress,
+        error=error
+    )
+    await websocket.send(progress_msg.to_json())
+
+async def load_model(model_path: str, lora_path: str = None) -> StableDiffusionXLImg2ImgPipeline:
+    """Load the model with progress reporting."""
+    try:
+        # Load base SDXL pipeline
+        pipe = StableDiffusionXLPipeline.from_single_file(
+            model_path,
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+            height=1024, width=1024,
+        ).to("cuda")
+        
+        # Convert to img2img pipeline
+        pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
+        
+        # Set scheduler
+        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+        
+        # Load LCM-LoRA weights if provided
+        if lora_path:
+            pipe.load_lora_weights(
+                "F:/ComfyUI_windows_portable_nvidia/ComfyUI_windows_portable/ComfyUI/models/loras/LCM_LoRA_Weights_SDXL.safetensors",
+                adapter_name="lcm"
+            )
+            pipe.load_lora_weights(lora_path, adapter_name="lora")
+            pipe.set_adapters(["lcm", "lora"], adapter_weights=[1.0, 0.8])
+            #pipe.set_adapters(["lcm"], adapter_weights=[1.0])
+        #pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+        return pipe
+    except Exception as e:
+        raise Exception(f"Error loading model: {str(e)}")
+
+#https://huggingface.co/blog/simple_sdxl_optimizations
+async def generate_image(
+    pipe: StableDiffusionXLImg2ImgPipeline,
+    request: GenerateRequest
+) -> str:
+    """Generate image and return base64 string."""
+    try:
+
+        image_stream = io.BytesIO(base64.b64decode(request.image))
+
+        init_image = Image.open(image_stream).convert('RGB')
+        init_image = transform(init_image)  # Now it's a torch tensor
+
+        # If you need it on GPU
+        init_image = init_image.to('cuda')  # Shape will be [3, H, W]
+
+        # If you need to match SDXL's expected format (batch dimension)
+        init_image = init_image.unsqueeze(0)  # Shape will be [1, 3, H, W]
+        prompt = "safe_pos, 1girl, solo, solo focus, 8k, masterpiece, hires, absurdres, splash art, gradient, looking at viewer, medium breasts, vibrant yellow hair, purple hair tips, (shoulder freckles:1.2), nude, athletic, confident, hand on waist, extremely long hair, dancing, delicate gold jewelry, thin gold bangles, belly button chain, action shot,, head tilt, hair censor, spotlight, tottotonero"
+        if len(request.prompt) == 0:
+            request.prompt = prompt
+
+        # Generate image
+        output_image = pipe(
+            request.prompt,
+            image=init_image,
+            num_inference_steps=request.num_inference_steps,
+            guidance_scale=request.guidance_scale,
+            strength=request.strength,
+            generator=torch.manual_seed(0)
+        ).images[0]
+        
+        # Convert to base64
+        buffered = BytesIO()
+        output_image.save(buffered, format="PNG")
+        buffered.seek(0) 
+        img_str = base64.b64encode(buffered.getvalue()).decode()
+        
+        return f"data:image/png;base64,{img_str}"
+    except Exception as e:
+        raise Exception(f"Error generating image: {str(e)} at line {e.__traceback__.tb_lineno}")
+
+async def handle_client(websocket):
+    """Handle WebSocket client connection."""
+    pipe = None
+    try:
+        async for message in websocket:
+            data = json.loads(message)
+            command = data.get("command")
+
+            if command == "load_model":
+                await send_progress(websocket, "Loading model...", 0.0)
+                request = LoadModelRequest.from_json(message)
+                
+                pipe = await load_model(request.model_path, request.lora_path)
+                await send_progress(websocket, "Model loaded successfully", 1.0)
+            elif command == "generate":
+                if not pipe:
+                    raise Exception("Model not loaded. Please load model first.")
+                
+                await send_progress(websocket, "Generating image...", 0.0)
+                request = GenerateRequest.from_json(message)
+                
+                result = await generate_image(pipe=pipe, request=request)
+                
+                response = GenerateResponse(image=result)
+                await websocket.send(response.to_json())
+            
+            else:
+                raise Exception(f"Unknown command: {command}")
+                
+    except Exception as e:
+        await send_progress(websocket, "Error", 0.0, str(e))
+
+async def server():
+    """Start WebSocket server."""
+    async with websockets.serve(handle_client, "localhost", 8765,max_size=None):
+        print("Server started on ws://localhost:8765")
+        await asyncio.Future()  # run forever
+
+if __name__ == "__main__":
+    # test()  # Comment out test function
+    asyncio.run(server())

frontend/realtime/paint_py_server/server.py

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

frontend/realtime/paint_py_server/test_copy.py

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

frontend/realtime/src/KonvaApp/Engine.ts

@@ -196,7 +196,6 @@ export class Engine {
         this.selectorSquare.disable(); // this breaks the paint if you use it in the wrong spot
         this.selectionManager.disable(); // this prevents selection.
         this.realTimeDrawEngine.enablePaintMode();
-        this.realTimeDrawEngine.disableDragging();
 
         uiAccess.toolbarMain.enable();
         uiAccess.toolbarMain.changeButtonState(ToolbarMainButtonNames.SELECT, {
@@ -208,7 +207,6 @@ export class Engine {
       case AppModes.SELECT: {
         console.log("APPMODE: SELECT");
         this.realTimeDrawEngine.disablePaintMode();
-        this.realTimeDrawEngine.enableDragging();
 
         this.selectorSquare.enable();
         this.selectionManager.enable();
@@ -466,12 +464,14 @@ export class Engine {
       }
     });
 
-    uiEvents.promptEvents.onPromptStrengthChanged((strength) => {
-      this.realTimeDrawEngine.currentStrength = strength;
+    uiEvents.promptEvents.onPromptStrengthChanged(async (strength) => {
+      this.realTimeDrawEngine.currentStrength = strength / 100.0;
+      // this.realTimeDrawEngine.render();
     });
 
-    uiEvents.promptEvents.onPromptTextChanged((prompt) => {
+    uiEvents.promptEvents.onPromptTextChanged(async (prompt) => {
       this.realTimeDrawEngine.currentPrompt = prompt;
+      // this.realTimeDrawEngine.render();
     });
 
     uiEvents.onAddShapeToEngine((shapeData) => {
@@ -652,7 +652,6 @@ export class Engine {
         await this.realTimeDrawEngine.render();
       },
     });
-    shapeNode.kNode.zIndex(1); // Replace desiredZIndex with a number
 
     this.commandManager.createNode(shapeNode);
 
@@ -682,7 +681,7 @@ export class Engine {
         await this.realTimeDrawEngine.render();
       },
     });
-    imageNode.kNode.zIndex(1);
+
     this.commandManager.createNode(imageNode);
   }
   public addPaintNode(

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

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

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

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

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

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

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

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

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

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

frontend/realtime/src/PyServer/DrawServer.ts

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

frontend/realtime/src/PyServer/DrawServerTypes.ts

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

frontend/realtime/src/PyServer/ServerClient.ts

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

frontend/realtime/src/PyServer/ServerTypes.ts

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

frontend/realtime/src/types/ServerTypes.ts

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

frontend/realtime/start_server.ps1

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