Feature/add background remove (#670)

Co-authored-by: CrossPr0duct <alicextensor@gmail.com>

crates/desktop/artcraft/src/core/state/data_dir/app_data_root.rs

@@ -33,6 +33,7 @@ pub struct AppDataRoot {
 impl AppDataRoot {
   pub fn create_default() -> anyhow::Result<Self> {
     let directory = get_default_data_dir()?;
+    println!("App data directory: {:?}", directory);
     Self::create_existing(directory)
   }
 

frontend/apps/editor3d/app/src/pages/PageDraw/Node.ts

@@ -21,6 +21,7 @@ export interface NodeProps {
   offsetX?: number;
   offsetY?: number;
   zIndex?: number;
+  locked?: boolean;
 }
 
 export class Node {
@@ -44,6 +45,7 @@ export class Node {
   offsetX: number;
   offsetY: number;
   zIndex: number;
+  locked: boolean;
   
   constructor({
     id,
@@ -65,7 +67,8 @@ export class Node {
     scaleY = 1,
     offsetX = 0,
     offsetY = 0,
-    zIndex = 0
+    zIndex = 0,
+    locked = false
   }: NodeProps) {
     this.id = id;
     this.x = x;
@@ -87,6 +90,7 @@ export class Node {
     this.offsetX = offsetX;
     this.offsetY = offsetY;
     this.zIndex = zIndex;
+    this.locked = locked;
   }
 
   setPosition(x: number, y: number): void {

frontend/apps/editor3d/app/src/pages/PageDraw/PageDraw.tsx

@@ -1,6 +1,7 @@
 import { useState, useRef, useEffect } from "react";
 import { PaintSurface } from "./PaintSurface";
 // https://github.com/SaladTechnologies/comfyui-api
+
 import "./App.css";
 import PromptEditor from "./PromptEditor/PromptEditor";
 import SideToolbar from "./components/ui/SideToolbar";
@@ -17,6 +18,43 @@ import Konva from "konva"; // just for types
 import { setCanvasRenderBitmap } from "../../signals/canvasRenderBitmap"
 import { captureStageImageBitmap } from "./hooks/useUpdateSnapshot"
 import { ContextMenuContainer } from "./components/ui/ContextMenu";
+import {
+  FalBackgroundRemoval,
+  FalBackgroundRemovalRequest,
+} from "@storyteller/tauri-api";
+
+
+  /**
+ * This decodes Base64 encoded PNG images into an image element.
+ * @param base64String a standard (not "web safe") base64-encoded string
+ */
+  export const DecodeBase64ToImage = async (base64String: string) : Promise<ImageBitmap> => {
+    // Create an image element
+    const img = document.createElement("img");
+  
+    // Convert base64 to data URL if it doesn't include the prefix
+    const dataUrl = base64String.startsWith("data:")
+      ? base64String
+      : `data:image/png;base64,${base64String}`;
+  
+    // Create a promise to handle the image loading
+    return new Promise((resolve, reject) => {
+      img.onload = async () => {
+        try {
+          const bitmap = await createImageBitmap(img);
+          resolve(bitmap);
+        } catch (error) {
+          reject(error);
+        }
+      };
+  
+      img.onerror = () => reject(new Error("Failed to load image"));
+  
+      // Set the source to trigger loading
+      img.src = dataUrl;
+    });
+  }
+
 const PageDraw = () => {
   //useStateSceneLoader();
 
@@ -273,7 +311,6 @@ const PageDraw = () => {
       <div className="relative z-0">
       <ContextMenuContainer 
         onAction={(e, action) => {
-
           if (action === "contextMenu") {
             const hasSelection = store.selectedNodeIds.length > 0;
             if (hasSelection) {
@@ -285,11 +322,40 @@ const PageDraw = () => {
               return false
             }
           }
-
+          return false
         }} 
-        onMenuAction={(action) => {
-        
+        onMenuAction={async (action) => {
           switch (action) {
+            case 'LOCK':
+              store.toggleLock(store.selectedNodeIds);
+              break;
+            case 'REMOVE_BACKGROUND':
+              // returns a success if we have selected images only
+              await store.removeBackground(store.selectedNodeIds, 
+                async (success: boolean, image_base64: string, message: string) => {
+                if (!success) {
+                  console.error(message);
+                  return { success: false };
+                }
+                try {
+                  const response = await FalBackgroundRemoval({ base64_image: image_base64 });
+                  if (response.status !== "success" || !("payload" in response)) {
+                    console.error("Failed to remove background", response);
+                    return { success: false };
+                  }
+
+                  const base64String = response.payload?.base64_bytes as string;
+                  const binaryString = atob(base64String);
+                  const bytes = Uint8Array.from(binaryString, c => c.charCodeAt(0));
+                  const blob = new Blob([bytes], { type: "image/png" });
+                  const file = new File([blob], "generated_image.png", { type: blob.type });
+                  return { success: true, file };
+                } catch (error) {
+                  console.error("Failed to remove background", error);
+                  return { success: false };
+                }
+              });
+              break;
             case 'BRING_TO_FRONT':
               store.bringToFront(store.selectedNodeIds);
               break;
@@ -302,10 +368,22 @@ const PageDraw = () => {
             case 'SEND_TO_BACK':
               store.sendToBack(store.selectedNodeIds);
               break;
+            case 'DUPLICATE':
+              store.copySelectedItems()
+              store.pasteItems()
+              break;
+            case 'DELETE':
+              store.deleteSelectedItems()
+              break;
             default:
               // No action needed for unhandled cases
           }
         }}
+        isLocked={store.selectedNodeIds.some(id => {
+          const node = store.nodes.find(n => n.id === id);
+          const lineNode = store.lineNodes.find(n => n.id === id);
+          return (node?.locked || lineNode?.locked) ?? false;
+        })}
       >
         <PaintSurface
           nodes={store.nodes}

frontend/apps/editor3d/app/src/pages/PageDraw/PaintSurface.tsx

@@ -546,6 +546,25 @@ export const PaintSurface = ({
         return;
       }
 
+      // allow selection
+      if (e.evt.button === 2) {
+        const node = store.nodes.find(n => n.id === nodeId);
+        const lineNode = store.lineNodes.find(n => n.id === nodeId);
+        const isLocked = (node?.locked || lineNode?.locked) ?? false;
+        
+        // If locked, only allow selection for context menu
+        if (isLocked) {
+          store.selectNode(nodeId);
+          return;
+        }
+      }
+      
+      // Don't select locked nodes
+      const node = nodes.find(n => n.id === nodeId) || store.lineNodes.find(n => n.id === nodeId);
+      if (node?.locked) {
+        return;
+      }
+
       // Check if Ctrl/Cmd key is pressed
       const isMultiSelect = e.evt.ctrlKey || e.evt.metaKey;
 
@@ -762,7 +781,7 @@ export const PaintSurface = ({
             onMouseLeave={handleNodeMouseLeave}
             draggable={draggableIfToolsNotActive(
               activeTool,
-              lineNode.draggable,
+              lineNode.draggable && lineNode.locked == false,
             )}
             onDragMove={(e) => handleLineDragMove(e, lineNode.id)}
             onDragStart={(e) => handleLineDragStart(e, lineNode.id)}
@@ -799,7 +818,7 @@ export const PaintSurface = ({
             offsetX={node.offsetX || 0}
             offsetY={node.offsetY || 0}
             zIndex={node.zIndex}
-            draggable={draggableIfToolsNotActive(activeTool, node.draggable)}
+            draggable={draggableIfToolsNotActive(activeTool, node.draggable && node.locked == false)}
             onMouseDown={(e) => handleNodeMouseDown(e, node.id)}
             onTap={(e) => handleNodeMouseDown(e, node.id)}
             onMouseEnter={handleNodeMouseEnter}
@@ -831,7 +850,7 @@ export const PaintSurface = ({
             offsetX={node.offsetX || 0}
             offsetY={node.offsetY || 0}
             zIndex={node.zIndex}
-            draggable={draggableIfToolsNotActive(activeTool, node.draggable)}
+            draggable={draggableIfToolsNotActive(activeTool, node.draggable && node.locked == false)}
             onMouseDown={(e) => handleNodeMouseDown(e, node.id)}
             onTap={(e) => handleNodeMouseDown(e, node.id)}
             onMouseEnter={handleNodeMouseEnter}
@@ -863,7 +882,7 @@ export const PaintSurface = ({
             offsetX={node.offsetX || 0}
             offsetY={node.offsetY || 0}
             zIndex={node.zIndex}
-            draggable={draggableIfToolsNotActive(activeTool, node.draggable)}
+            draggable={draggableIfToolsNotActive(activeTool, node.draggable && node.locked == false)}
             onMouseDown={(e) => handleNodeMouseDown(e, node.id)}
             onTap={(e) => handleNodeMouseDown(e, node.id)}
             onMouseEnter={handleNodeMouseEnter}
@@ -911,7 +930,7 @@ export const PaintSurface = ({
             scaleY={node.scaleY || 1}
             offsetX={node.offsetX || 0}
             offsetY={node.offsetY || 0}
-            draggable={draggableIfToolsNotActive(activeTool, node.draggable)}
+            draggable={draggableIfToolsNotActive(activeTool, node.draggable && node.locked == false)}
             onMouseDown={(e) => handleNodeMouseDown(e, node.id)}
             onTap={(e) => handleNodeMouseDown(e, node.id)}
             onMouseEnter={handleNodeMouseEnter}

frontend/apps/editor3d/app/src/pages/PageDraw/components/ui/ContextMenu.tsx

@@ -67,8 +67,9 @@ interface ContextMenuContainerProps {
   children: React.ReactNode;
   items?: MenuItem[];
   onAction?: (e: React.MouseEvent, actionName: string) => boolean | void;
-  onMenuAction?: (action: string) => void;  // New prop for menu item actions
+  onMenuAction?: (action: string) => void;
   className?: string;
+  isLocked?: boolean;
 }
 
 export const ContextMenuContainer: React.FC<ContextMenuContainerProps> = ({ 
@@ -76,11 +77,27 @@ export const ContextMenuContainer: React.FC<ContextMenuContainerProps> = ({
   items = defaultMenuItems,
   onAction,
   onMenuAction,
-  className = ''
+  className = '',
+  isLocked = false
 }) => {
   const [showMenu, setShowMenu] = React.useState(false);
   const [position, setPosition] = React.useState<MenuPosition>({ x: 0, y: 0 });
 
+  const filteredItems = React.useMemo(() => {
+    if (isLocked) {
+      return [
+        {
+          icon: <svg className="w-4 h-4" viewBox="0 0 24 24" fill="none" stroke="currentColor">
+            <path d="M8 11V7a4 4 0 118 0m-4 8v2m-6 4h12a2 2 0 002-2v-6a2 2 0 00-2-2H6a2 2 0 00-2 2v6a2 2 0 002 2z" strokeWidth="2" strokeLinecap="round"/>
+          </svg>,
+          label: 'Unlock',
+          action: 'LOCK'
+        }
+      ];
+    }
+    return items;
+  }, [isLocked, items]);
+
   const handleContextMenu = (e: React.MouseEvent) => {
     e.preventDefault();
     const newPosition = { x: e.clientX, y: e.clientY };
@@ -111,7 +128,7 @@ export const ContextMenuContainer: React.FC<ContextMenuContainerProps> = ({
       {children}
       {showMenu && (
         <ContextMenu
-          items={items}
+          items={filteredItems}
           onAction={handleMenuAction}
           position={position}
           onClose={() => setShowMenu(false)}
@@ -174,5 +191,24 @@ const defaultMenuItems: MenuItem[] = [
     label: 'Duplicate',
     action: 'DUPLICATE'
   },
+  {
+    icon: (
+      <svg
+        className="w-4 h-4"
+        viewBox="0 0 24 24"
+        fill="none"
+        stroke="currentColor"
+      >
+        <path
+          d="M19 7l-.867 12.142A2 2 0 0116.138 21H7.862a2 2 0 01-1.995-1.858L5 7m5 4v6m4-6v6M9 7h6m-7 0V4a1 1 0 011-1h4a1 1 0 011 1v3"
+          strokeWidth="2"
+          strokeLinecap="round"
+          strokeLinejoin="round"
+        />
+      </svg>
+    ),
+    label: 'Delete',
+    action: 'DELETE'
+  },
 ];
 

frontend/apps/editor3d/app/src/pages/PageDraw/stores/SceneState.ts

@@ -18,6 +18,29 @@ export type LineNode = {
   offsetX?: number;
   offsetY?: number;
   zIndex: number;
+  locked?: boolean; // Add locked property
+};
+
+
+   // Logic to remove background from image nodes would go here
+const convertFileToBase64 = (file: File): Promise<string> => {
+    return new Promise((resolve, reject) => {
+        const reader = new FileReader();
+        
+        reader.onloadend = () => {
+            if (reader.result) {
+                resolve(reader.result as string);
+            } else {
+                reject(new Error("Failed to convert file to base64."));
+            }
+        };
+        
+        reader.onerror = () => {
+            reject(new Error("Error reading file."));
+        };
+        
+        reader.readAsDataURL(file);
+    });
 };
 
 export type ActiveTool = 'select' | 'draw' | 'eraser' | 'backgroundColor' | 'shape';
@@ -110,6 +133,10 @@ interface SceneState {
   sendToBack: (nodeIds: string[]) => void;
   bringForward: (nodeIds: string[]) => void;
   sendBackward: (nodeIds: string[]) => void;
+  removeBackground: (nodeIds: string[], operation: (success: boolean,base64_image:string, message: string) => Promise<{ success: boolean; file?: File }>) => Promise<void>;
+
+  // Add new lock action
+  toggleLock: (nodeIds: string[]) => void;
 }
 
 const generateId = (): string => {
@@ -1131,4 +1158,90 @@ export const useSceneStore = create<SceneState>((set, get) => ({
     });
     get().saveState();
   },
+  removeBackground: async (nodeIds: string[], operation:  (success: boolean, base64_image:string,message: string) => Promise<{ success: boolean; file?: File }>) => {
+   
+    const hasImageNodes = nodeIds.some(id => {
+        const node = get().nodes.find(n => n.id === id);
+        return node ? node.type === 'image': false;
+    });
+    let callbackSuccess = false
+    if (hasImageNodes) {
+        console.log("Selected nodes include image types.");
+     
+        const firstNode = get().nodes.find(node => nodeIds.includes(node.id) && node.type === 'image');
+        if (firstNode) {
+            const imageFile = firstNode.imageFile; // Assuming the file is stored in the node
+            if (imageFile !== undefined) {
+              try {
+                const base64Image = await convertFileToBase64(imageFile);
+                console.log("Base64 Image:", base64Image);
+                // Additional logic to handle the base64 image can be added here
+                const { success, file } = await operation(true, base64Image, "Removing Background."); // Indicate success in removing background
+                callbackSuccess = success;
+                if (callbackSuccess) {
+                  // save the state of the image.
+                  set((state: SceneState) => ({
+                    nodes: state.nodes.map((node: Node) => {
+                      if (node.id === firstNode.id) {
+                        // Create an object URL for immediate preview if needed
+                        console.log("Updated Node!")
+                        const updatedImageUrl: string | undefined = file ? URL.createObjectURL(file) : node.imageUrl;
+                        return new Node({
+                          ...node,
+                          imageFile: file,
+                          imageUrl: updatedImageUrl,
+                        });
+                      }
+                      return node;
+                    }),
+                  }));
+                
+                }
+              } catch (error) {
+                console.error("Error converting file to base64:", error);
+              }
+            }
+            console.log("Image File:", imageFile);
+        } else {
+            console.warn("No image node found or first node is not an image.");
+        }
+    } else {
+        console.log("No selected nodes are of image type.");
+        operation(false,"","Did not select a Images"); // Indicate failure to remove background
+    }
+  
+  },
+
+  // Add toggleLock action
+  toggleLock: (nodeIds: string[]) => {
+    set((state) => {
+      // Update regular nodes
+      const updatedNodes = state.nodes.map(node => {
+        if (nodeIds.includes(node.id)) {
+          return new Node({
+            ...node,
+            locked: !node.locked
+          });
+        }
+        return node;
+      });
+
+      // Update line nodes
+      const updatedLineNodes = state.lineNodes.map(node => {
+        if (nodeIds.includes(node.id)) {
+          return {
+            ...node,
+            locked: !node.locked
+          };
+        }
+        return node;
+      });
+
+      return {
+        nodes: updatedNodes,
+        lineNodes: updatedLineNodes
+      };
+    });
+    get().saveState();
+  }
 }));