render bug fix

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
                     

frontend/realtime/src/KonvaApp/Engine.ts

@@ -3,7 +3,7 @@ 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 {
@@ -129,6 +129,9 @@ export class Engine {
       onDraw: async (canvas, lineBounds) => {
         await this.addPaintNode(canvas, lineBounds);
       },
+      onPreviewCopy: async (image) => {
+        await this.addPreviewCopy(image);
+      },
     });
 
     // Selector Square to select Nodes
@@ -652,7 +655,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: 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,

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

@@ -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.
 
@@ -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();
         }
       });
     });
@@ -774,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:")
@@ -854,26 +854,26 @@ export class RealTimeDrawEngine {
       return;
     }
 
-    //   try {
-    //     const base64Bitmap = await this.imageBitmapToBase64(bitmap);
+    try {
+      const base64Bitmap = await this.imageBitmapToBase64(bitmap);
 
-    //     const base64BitmapResponse = await invoke("infer_image", {
-    //       image: base64Bitmap,
-    //       prompt: this.currentPrompt,
-    //       strength: this.currentStrength,
-    //     });
+      const base64BitmapResponse = await invoke("infer_image", {
+        image: base64Bitmap,
+        prompt: this.currentPrompt,
+        strength: this.currentStrength,
+      });
 
-    //     console.log(base64BitmapResponse);
-    //     const decoded = await this.base64ToImageBitmap(
-    //       base64BitmapResponse as string,
-    //     );
+      console.log(base64BitmapResponse);
+      const decoded = await this.base64ToImageBitmap(
+        base64BitmapResponse as string,
+      );
 
-    //     this.outputBitmap = decoded;
-    //     this.previewCanvas.image(decoded);
-    //   } catch (error) {
-    //     console.error("Error during image processing:", error);
-    //   } finally {
-    //     this.isProcessing = false;
-    //   }
+      this.outputBitmap = decoded;
+      this.previewCanvas.image(decoded);
+    } catch (error) {
+      console.error("Error during image processing:", error);
+    } finally {
+      this.isProcessing = false;
+    }
   }
 }

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

@@ -5,7 +5,6 @@ import {
   faEraser,
   faFilePlus,
   faImage,
-  faImages,
   faLocationArrow,
   faPaintbrush,
   faShapes,
@@ -28,8 +27,6 @@ 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";
-import { BackgroundMenu } from "~/components/ui/ToolbarMain/BackgroundMenu";
-import { bgColor } from "~/signals/uiEvents/toolbarMain/backgroundMenu";
 
 export const ToolbarMain = ({
   disabled = false,
@@ -229,31 +226,6 @@ export const ToolbarMain = ({
             tooltip="Save"
           /> */}
         </div>
-
-        <hr className="w-full border-t border-white/15" />
-
-
-        <div className="flex flex-col items-center gap-2">
-          <div className="relative">
-            <BackgroundMenu
-              color={bgColor.value}
-              onChange={dispatchUiEvents.toolbarMain.setBgColor}
-              faIcon={faImage}
-              borderStyle="bg-transparent hover:bg-white/15 text-white"
-              showBar={false}
-              staticIconColor="white"
-              streamChanges
-              defaultOpen={false}
-              anchor="right"
-              anchorGap={12}
-              closeOnMouseLeave={true}
-            />
-            <div
-              className="pointer-events-none absolute -bottom-1.5 -right-1.5 h-[18px] w-[18px] rounded-full border border-gray-400"
-              style={{ backgroundColor: bgColor.value }}
-            />
-          </div>
-        </div>
       </div>
     </>
   );

frontend/realtime/src/components/features/ToolbarMain/enum.ts

@@ -15,5 +15,4 @@ export enum ToolbarMainButtonNames {
   SAVE = "SAVE",
   PREVIEW = "PREVIEW",
   DOWNLOAD = "DOWNLOAD",
-  BACKGROUND = "BACKGROUND",
 }

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

@@ -1,180 +0,0 @@
-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";
-
-export const BackgroundMenu = ({
-  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();
-
-  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
-              >
-                <HexColorPicker
-                  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}
-                />
-                {!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/backgroundMenu.ts

@@ -1,15 +0,0 @@
-import { signal, effect } from "@preact/signals-react";
-
-export const DEFAULT_BG_COLOR: string = "#ffffff";
-export const bgColor = signal<string>(DEFAULT_BG_COLOR);
-export const setBgColor = (data: string) => {
-  bgColor.value = data;
-};
-
-export const onBgColorChanged = (callback: (data: string) => void) => {
-  effect(() => {
-    if (bgColor.value) {
-      callback(bgColor.value);
-    }
-  });
-};

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

@@ -5,19 +5,16 @@ import {
 } from "./loadingBar";
 import { onEraseBrushSizeChanged, setEraseBrushSize } from "./eraseMode";
 import { onPaintColorChanged, setPaintBrushSize, setPaintColor } from "./paintMode";
-import { onBgColorChanged, setBgColor } from "./backgroundMenu";
 export const dispatchers = {
   ...buttonDispatchers,
   loadingBarRetry: loadingBarRetryDispatch,
   setPaintColor,
   setPaintBrushSize,
-  setEraseBrushSize,
-  setBgColor
+  setEraseBrushSize
 };
 export const eventsHandlers = {
   ...buttonEventsHandlers,
   loadingBarRetry: loadingBarRetryEventHandler,
   onPaintColorChanged,
-  onEraseBrushSizeChanged,
-  onBgColorChanged
+  onEraseBrushSizeChanged
 };