Add Background button and menu to toolbar

frontend/realtime/paint_py_server/server.py

@@ -10,13 +10,7 @@ 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 diffusers import AutoPipelineForImage2Image, LCMScheduler, UNet2DConditionModel
 from tqdm import tqdm
 from huggingface_hub import HfApi, hf_hub_download
 
@@ -99,7 +93,6 @@ 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:
@@ -109,11 +102,10 @@ async def load_model(sdxl_checkpoint_path: pathlib.Path, progress_callback: Opti
         
         # Load the base model using the checkpoint path
         pipe = AutoPipelineForImage2Image.from_pretrained(
-            "stabilityai/stable-diffusion-xl-base-1.0",
+            sdxl_checkpoint_path,
             unet=unet,
             torch_dtype=torch.float16,
             variant="fp16",
-           use_safetensors=True
         ).to("cuda")
         
         if progress_callback:
@@ -129,7 +121,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",from_safetensors=True)
+        pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl", adapter_name="lcm")
         
         if progress_callback:
             await progress_callback(100, 100)
@@ -236,80 +228,36 @@ async def handle_client(websocket):
         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)}")
+                logger.info(f"Received message: {data.keys()}")
                 
                 # 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')
+                    setup_data = data.get('setup', {})
                     
-                    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")
+                    # 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
                     
                     # 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
+                        'lora_path': str(lora_path) if lora_path else None,
+                        # Add other settings as needed
                     }
-                    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
+                        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:
@@ -322,13 +270,13 @@ async def handle_client(websocket):
                     
                     # 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 lora_path and lora_path.exists():
+                        lora_success = await load_lora(pipe, lora_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"])}'
+                                'error': f'Failed to load LoRA from {lora_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,9 +129,6 @@ export class Engine {
       onDraw: async (canvas, lineBounds) => {
         await this.addPaintNode(canvas, lineBounds);
       },
-      onPreviewCopy: async (image) => {
-        await this.addPreviewCopy(image);
-      },
     });
 
     // Selector Square to select Nodes
@@ -655,18 +652,7 @@ 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

@@ -1,58 +0,0 @@
-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,8 +73,6 @@ export class RealTimeDrawEngine {
     },
   ) => void;
 
-  private onPreviewCopyCallback?: (previewCopy: Konva.Image) => void; // New Callback
-
   private serverSocket: WebSocket | null = null;
 
   constructor({
@@ -83,7 +81,6 @@ export class RealTimeDrawEngine {
     mediaLayerRef,
     offScreenCanvas,
     onDraw,
-    onPreviewCopy, // New Parameter
   }: {
     width: number;
     height: number;
@@ -98,13 +95,11 @@ 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.
 
@@ -208,7 +203,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}%`);
@@ -431,6 +426,16 @@ 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(),
@@ -442,20 +447,15 @@ export class RealTimeDrawEngine {
         listening: true,
       });
 
-      // 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);
-      }
+      this.drawingsLayer.add(previewCopy); // Add to drawingsLayer instead of mediaLayerRef
+      this.drawingsLayer.batchDraw();
 
       // 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 = document.createElement("img");
+    const img = new Image();
 
     // 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,6 +5,7 @@ import {
   faEraser,
   faFilePlus,
   faImage,
+  faImages,
   faLocationArrow,
   faPaintbrush,
   faShapes,
@@ -27,6 +28,8 @@ 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,
@@ -226,6 +229,31 @@ 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,4 +15,5 @@ export enum ToolbarMainButtonNames {
   SAVE = "SAVE",
   PREVIEW = "PREVIEW",
   DOWNLOAD = "DOWNLOAD",
+  BACKGROUND = "BACKGROUND",
 }

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

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

@@ -0,0 +1,15 @@
+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,16 +5,19 @@ 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
+  setEraseBrushSize,
+  setBgColor
 };
 export const eventsHandlers = {
   ...buttonEventsHandlers,
   loadingBarRetry: loadingBarRetryEventHandler,
   onPaintColorChanged,
-  onEraseBrushSizeChanged
+  onEraseBrushSizeChanged,
+  onBgColorChanged
 };