From 8fea9c1a5e784607db8ddf35659bc2bf34be2c6e Mon Sep 17 00:00:00 2001 From: CrossPr0duct Date: Tue, 25 Feb 2025 23:09:48 -0500 Subject: [PATCH] Working on previews --- frontend/realtime/paint_py_server/server.py | 90 +++++++++++++++---- frontend/realtime/src/KonvaApp/Engine.ts | 13 ++- .../src/KonvaApp/Nodes/PreviewCopy.ts | 58 ++++++++++++ .../RenderingPrimitives/RealTimeDrawEngine.ts | 36 ++++---- 4 files changed, 159 insertions(+), 38 deletions(-) create mode 100644 frontend/realtime/src/KonvaApp/Nodes/PreviewCopy.ts diff --git a/frontend/realtime/paint_py_server/server.py b/frontend/realtime/paint_py_server/server.py index e7a464c371..2ed7f16767 100644 --- a/frontend/realtime/paint_py_server/server.py +++ b/frontend/realtime/paint_py_server/server.py @@ -10,7 +10,13 @@ from typing import Dict, Any, Optional, Tuple, Callable, List import torch import websockets from PIL import Image -from diffusers import AutoPipelineForImage2Image, LCMScheduler, UNet2DConditionModel +from diffusers import ( + StableDiffusionXLPipeline, + StableDiffusionXLImg2ImgPipeline, + AutoPipelineForImage2Image, + LCMScheduler, + UNet2DConditionModel, +) from tqdm import tqdm from huggingface_hub import HfApi, hf_hub_download @@ -93,6 +99,7 @@ async def load_model(sdxl_checkpoint_path: pathlib.Path, progress_callback: Opti "latent-consistency/lcm-sdxl", torch_dtype=torch.float16, variant="fp16", + use_safetensors=True, ) if progress_callback: @@ -102,10 +109,11 @@ async def load_model(sdxl_checkpoint_path: pathlib.Path, progress_callback: Opti # Load the base model using the checkpoint path pipe = AutoPipelineForImage2Image.from_pretrained( - sdxl_checkpoint_path, + "stabilityai/stable-diffusion-xl-base-1.0", unet=unet, torch_dtype=torch.float16, variant="fp16", + use_safetensors=True ).to("cuda") if progress_callback: @@ -121,7 +129,7 @@ async def load_model(sdxl_checkpoint_path: pathlib.Path, progress_callback: Opti await progress_callback(85, 100) # Load the default LCM adapter - pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl", adapter_name="lcm") + pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl", adapter_name="lcm",from_safetensors=True) if progress_callback: await progress_callback(100, 100) @@ -228,36 +236,80 @@ async def handle_client(websocket): async for message in websocket: try: data = json.loads(message) - logger.info(f"Received message: {data.keys()}") + 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': - setup_data = data.get('setup', {}) + logger.info("Processing setup message") + model_data = data.get('model', {}) + lora_data = data.get('lora') + device = data.get('device', 'cuda') - # Extract setup parameters - sdxl_checkpoint_path = pathlib.Path(setup_data.get('sdxl_checkpoint_path', '')) - lora_path = pathlib.Path(setup_data.get('lora_path', '')) if setup_data.get('lora_path') else None + 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_path) if lora_path else None, - # Add other settings as needed + '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(): - await websocket.send(json.dumps({ - 'type': 'setup_response', - 'success': False, - 'error': f'SDXL checkpoint path does not exist: {sdxl_checkpoint_path}' - })) - continue + 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: @@ -270,13 +322,13 @@ async def handle_client(websocket): # Load LoRA if specified lora_success = True - if lora_path and lora_path.exists(): - lora_success = await load_lora(pipe, lora_path, progress_callback=progress_callback) + 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 {lora_path}' + 'error': f'Failed to load LoRA from {pathlib.Path(lora_data["path"])}' })) continue diff --git a/frontend/realtime/src/KonvaApp/Engine.ts b/frontend/realtime/src/KonvaApp/Engine.ts index 697f4a7d64..3f90d809de 100644 --- a/frontend/realtime/src/KonvaApp/Engine.ts +++ b/frontend/realtime/src/KonvaApp/Engine.ts @@ -652,7 +652,18 @@ export class Engine { console.debug("Added shapenode:", shapeNode); console.debug("Added node's ID:", shapeNode.kNode._id); } - + public addPreviewCopy(image: Image.Konva) { + // Start of Selection + const copyNode = new PreviewCopyNode({ + image: previewCopy, + mediaLayerRef: this.mediaLayerRef, + selectionManagerRef: this.selectionManagerRef, + loaded: async () => { + this.drawingsLayer.draw(); + }, + }); + this.commandManager.createNode(imageNode); + } public addImage(imageFile: File) { const imageNode = new ImageNode({ mediaLayerRef: this.mediaLayer, diff --git a/frontend/realtime/src/KonvaApp/Nodes/PreviewCopy.ts b/frontend/realtime/src/KonvaApp/Nodes/PreviewCopy.ts new file mode 100644 index 0000000000..7aaa0afbb2 --- /dev/null +++ b/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; +} + +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; + } +} diff --git a/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts b/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts index bf3cea3789..b7603747fa 100644 --- a/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts +++ b/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts @@ -73,6 +73,8 @@ export class RealTimeDrawEngine { }, ) => void; + private onPreviewCopyCallback?: (previewCopy: Konva.Image) => void; // New Callback + private serverSocket: WebSocket | null = null; constructor({ @@ -81,6 +83,7 @@ export class RealTimeDrawEngine { mediaLayerRef, offScreenCanvas, onDraw, + onPreviewCopy, // New Parameter }: { width: number; height: number; @@ -95,11 +98,13 @@ export class RealTimeDrawEngine { 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. @@ -165,7 +170,7 @@ export class RealTimeDrawEngine { this.captureCanvas.setZIndex(0); this.previewCanvas.setZIndex(1); // Add mouse events for preview canvas copying - //this.previewCopyListener(); + this.previewCopyListener(); //this.startServer(); } @@ -203,7 +208,7 @@ export class RealTimeDrawEngine { 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}%`); @@ -426,16 +431,6 @@ export class RealTimeDrawEngine { 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({ x: this.previewCanvas.x(), @@ -447,15 +442,20 @@ export class RealTimeDrawEngine { listening: true, }); - this.drawingsLayer.add(previewCopy); // Add to drawingsLayer instead of mediaLayerRef - this.drawingsLayer.batchDraw(); + // 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.drawingsLayer.draw(); + // this.drawingsLayer.draw(); }); previewCopy.on("dragend", () => { @@ -468,11 +468,11 @@ export class RealTimeDrawEngine { x: this.captureCanvas.x(), y: this.captureCanvas.y(), }); - this.drawingsLayer.batchDraw(); + //this.drawingsLayer.batchDraw(); } else { // Remove if dropped outside capture area - previewCopy.destroy(); - this.drawingsLayer.batchDraw(); + //previewCopy.destroy(); + //this.drawingsLayer.batchDraw(); } }); });