From f37011ffc637f65cf48cc0ead86a10c6662bc234 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 20 Mar 2025 01:41:43 -0400 Subject: [PATCH 1/6] remove background: frontend UX (#458) --- .../src-tauri/src/endpoints/image_endpoint.rs | 1 + .../endpoints/remove_background_endpoint.rs | 3 +- frontend/realtime/src/KonvaApp/Engine.ts | 91 +++++++++++++++---- .../EngineUtitlities/CommandManager.ts | 14 +++ .../realtime/src/KonvaApp/Nodes/ImageNode.ts | 44 +++++++++ .../KonvaApp/NodesManagers/NodesManager.ts | 4 + .../NodesManagers/ToolbarNodeUtilities.ts | 12 ++- .../RenderingPrimitives/RealTimeDrawEngine.ts | 34 +------ .../components/features/ToolbarNode/data.ts | 5 + .../components/features/ToolbarNode/enums.ts | 1 + .../src/utilities/DecodeBase64ToImage.ts | 30 ++++++ .../utilities/EncodeImageBitmapToBase64.ts | 22 +++++ 12 files changed, 206 insertions(+), 55 deletions(-) create mode 100644 frontend/realtime/src/utilities/DecodeBase64ToImage.ts create mode 100644 frontend/realtime/src/utilities/EncodeImageBitmapToBase64.ts diff --git a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs index fa0db4e540..9709736ba0 100644 --- a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs @@ -32,6 +32,7 @@ pub async fn infer_image( app_data_root: State<'_, AppDataRoot>, app: AppHandle, ) -> Result { + info!("infer_image endpoint called."); let prompt = get_prompt_or_fallback(prompt); diff --git a/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs index 8d8b7aa5c9..487315b866 100644 --- a/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs @@ -6,7 +6,7 @@ use crate::state::app_dir::AppDataRoot; use crate::utils::image::decode_base64_image::decode_base64_image; use crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png; use image::{imageops, DynamicImage}; -use log::error; +use log::{error, info}; use ndarray::{Array, CowArray}; use ort::Value; use tauri::{AppHandle, State}; @@ -19,6 +19,7 @@ pub async fn remove_background( app_data_root: State<'_, AppDataRoot>, app: AppHandle, ) -> Result { + info!("remove_background endpoint called."); let image = decode_base64_image(image) .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?; diff --git a/frontend/realtime/src/KonvaApp/Engine.ts b/frontend/realtime/src/KonvaApp/Engine.ts index ef2814b226..91b3cbdb42 100644 --- a/frontend/realtime/src/KonvaApp/Engine.ts +++ b/frontend/realtime/src/KonvaApp/Engine.ts @@ -1,35 +1,36 @@ 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 {uiAccess, uiEvents} from "~/signals"; +import {ImageNode, ShapeNode, ShapeType, TextNode, VideoNode} from "./Nodes"; +import {PaintNode} from "./Nodes/PaintNode"; +import {PreviewCopyNode} from "./Nodes/PreviewCopy"; +import {UndoStackManager} from "./UndoRedo"; +import {CommandManager, MatteBox, SceneManager} from "./EngineUtitlities"; import { - NodesManager, NodeIsolator, - NodeTransformer, + NodesManager, NodesTranslationEventDetails, NodeTransformationEventDetails, + NodeTransformer, SelectionManager, SelectionManagerEvents, SelectorSquare, } from "./NodesManagers"; -import { ImageNode, VideoNode, TextNode } from "./Nodes"; -import { EngineOptions, TextNodeData, VideoNodeData } from "./types"; +import {EngineOptions, TextNodeData, VideoNodeData} from "./types"; -import { AppModes, VideoResolutions } from "./constants"; -import { ToolbarMainButtonNames } from "~/components/features/ToolbarMain/enum"; +import {AppModes, VideoResolutions} from "./constants"; +import {ToolbarMainButtonNames} from "~/components/features/ToolbarMain/enum"; -import { ToolbarNodeButtonNames } from "~/components/features/ToolbarNode/enums"; -import { NavigateFunction } from "react-router-dom"; -import { LoadingVideosProvider } from "./EngineUtitlities/LoadingVideosProvider"; +import {ToolbarNodeButtonNames} from "~/components/features/ToolbarNode/enums"; +import {NavigateFunction} from "react-router-dom"; +import {LoadingVideosProvider} from "./EngineUtitlities/LoadingVideosProvider"; -import { VideoExtractionHandler } from "./EngineUtitlities/VideoExtractionHandler/VideoExtractionHandler"; -import { RealTimeDrawEngine } from "./RenderingPrimitives/RealTimeDrawEngine"; -import { NodeColor } from "~/signals/uiEvents/toolbarNode"; -import { ShapeType } from "./Nodes"; +import {VideoExtractionHandler} from "./EngineUtitlities/VideoExtractionHandler/VideoExtractionHandler"; +import {RealTimeDrawEngine} from "./RenderingPrimitives/RealTimeDrawEngine"; +import {NodeColor} from "~/signals/uiEvents/toolbarNode"; +import {invoke} from "@tauri-apps/api/core"; +import {DecodeBase64ToImage} from "~/utilities/DecodeBase64ToImage.ts"; +import {EncodeImageBitmapToBase64} from "~/utilities/EncodeImageBitmapToBase64.ts"; export interface RenderingOptions { @@ -326,6 +327,9 @@ export class Engine { uiEvents.toolbarNode.MOVE_LAYER_UP.onClick(() => this.commandManager.moveNodesUp(), ); + uiEvents.toolbarNode.REMOVE_BACKGROUND.onClick(() => + this.removeBackground(), + ); // Listen to Toolbar Main uiEvents.toolbarMain.UNDO.onClick(() => { @@ -445,6 +449,38 @@ export class Engine { }); } + removeBackground() { + const selectedNodeSet = this.selectionManager.getSelectedNodes(); + + if (selectedNodeSet.size !== 1) { + return; // Can only remove background from a single image at a time. + } + + const [ selectedNode ] = selectedNodeSet; + + const removeBgInvocation + = async (selectedNode: ImageNode, callback: (existingImageNode: ImageNode, createdImage: ImageBitmap) => void) => { + + const image : HTMLImageElement = selectedNode.kNode.image(); + const base64Bitmap = await EncodeImageBitmapToBase64(image); + const base64BitmapResponse = await invoke("remove_background", { + image: base64Bitmap, + }); + + const createdImage = await DecodeBase64ToImage(base64BitmapResponse as string); + + callback(selectedNode, createdImage); + }; + + const _promise = removeBgInvocation(selectedNode, (selectedNode: ImageNode, createdImage: ImageBitmap) => { + if (!this.nodesManager.hasNode(selectedNode)) { + return; + } + this.commandManager.deleteSpecificNodes([selectedNode]); + this.addImageFromImageBitmap(createdImage); + }) + } + disableAllButtons() { const buttonNames = Object.values(ToolbarNodeButtonNames); for (const name of buttonNames) { @@ -624,6 +660,23 @@ export class Engine { this.commandManager.createNode(imageNode); } + + public addImageFromImageBitmap(imageBitmap: ImageBitmap) { + const imageNode = new ImageNode({ + mediaLayerRef: this.mediaLayer, + canvasPosition: this.realTimeDrawEngine.captureCanvas.position(), + canvasSize: this.realTimeDrawEngine.captureCanvas.size(), + imageBitmap: imageBitmap, + selectionManagerRef: this.selectionManager, + loaded: async () => { + await this.realTimeDrawEngine.render(); + this.setAppMode(AppModes.SELECT); + }, + }); + + this.commandManager.createNode(imageNode); + } + public addPaintNode( canvas: HTMLCanvasElement, lineBounds: { diff --git a/frontend/realtime/src/KonvaApp/EngineUtitlities/CommandManager.ts b/frontend/realtime/src/KonvaApp/EngineUtitlities/CommandManager.ts index 40aa02b087..5463d4e9e4 100644 --- a/frontend/realtime/src/KonvaApp/EngineUtitlities/CommandManager.ts +++ b/frontend/realtime/src/KonvaApp/EngineUtitlities/CommandManager.ts @@ -78,6 +78,20 @@ export class CommandManager { this.undoStackManagerRef.executeCommand(command); this.renderEngineRef.render(); } + + deleteSpecificNodes(nodes: MediaNode[]) { + const command = new DeleteCommand({ + nodes: nodes, + mediaLayerRef: this.mediaLayerRef, + nodesManagerRef: this.nodesManagerRef, + nodeTransformerRef: this.nodeTransformerRef, + selectionManagerRef: this.selectionManagerRef, + renderEngineRef: this.renderEngineRef, + }); + this.undoStackManagerRef.executeCommand(command); + this.renderEngineRef.render(); + } + toggleLockNodes() { const nodes = this.selectionManagerRef.getSelectedNodes(); const node = nodes.values().next().value; diff --git a/frontend/realtime/src/KonvaApp/Nodes/ImageNode.ts b/frontend/realtime/src/KonvaApp/Nodes/ImageNode.ts index fc57c971c2..a81cdf7301 100644 --- a/frontend/realtime/src/KonvaApp/Nodes/ImageNode.ts +++ b/frontend/realtime/src/KonvaApp/Nodes/ImageNode.ts @@ -12,6 +12,7 @@ interface ImageNodeContructor { canvasPosition: Position; canvasSize: Size; imageFile?: File; + imageBitmap?: ImageBitmap; mediaFileToken?: string; mediaFileUrl?: string; transform?: TransformationData; @@ -34,6 +35,7 @@ export class ImageNode extends NetworkedNode { mediaLayerRef, selectionManagerRef, loaded, + imageBitmap, }: ImageNodeContructor) { // kNodes need to be created first to guaruntee // that it is not undefined in parent's context @@ -73,6 +75,14 @@ export class ImageNode extends NetworkedNode { this.loadImageFromUrl(mediaFileUrl); return; } + if (imageBitmap) { + this.loadImageFromImageBitmap({ + imageBitmap: imageBitmap, + maxSize: canvasSize, + refPosition: canvasPosition, + }); + return; + } console.log("image node creation is fucked"); } private loadImageFromFile({ @@ -136,6 +146,40 @@ export class ImageNode extends NetworkedNode { this.listenToBaseKNode(); } + private loadImageFromImageBitmap({ + imageBitmap, + maxSize, + refPosition, + }: { + imageBitmap: ImageBitmap; + maxSize: Size; + refPosition: Position; + }) { + this.imageSize = { + width: imageBitmap.width, + height: imageBitmap.height, + }; + const adjustedSize = NodeUtilities.adjustNodeSizeToCanvas({ + componentSize: this.imageSize, + maxSize: maxSize, // TODO + }); + const centerPosition = NodeUtilities.positionNodeOnCanvasCenter({ + canvasOffset: refPosition, // TODO + componentSize: adjustedSize, // TODO + maxSize: maxSize, // TODO + }); + this.kNode.image(imageBitmap); + this.kNode.setSize(adjustedSize); + this.kNode.setPosition(centerPosition); + this.kNode.moveToTop(); + this.kNode.fill(transparent); + this.listenToBaseKNode(); + this.mediaLayerRef.draw(); + //this.uploadImage(imageFile); + this.loaded(); + this.listenToBaseKNode(); + } + private async uploadImage(imageFile: File) { //this.setProgress({ progress: 10, status: UploadStatus.UPLOADING }); diff --git a/frontend/realtime/src/KonvaApp/NodesManagers/NodesManager.ts b/frontend/realtime/src/KonvaApp/NodesManagers/NodesManager.ts index 5fec9128fd..f86afc2018 100644 --- a/frontend/realtime/src/KonvaApp/NodesManagers/NodesManager.ts +++ b/frontend/realtime/src/KonvaApp/NodesManagers/NodesManager.ts @@ -44,6 +44,10 @@ export class NodesManager { // this.printNodesAndZIndices(); } + public hasNode(node: MediaNode): boolean { + return this.allNodes.has(node); + } + /************************* * For Debugging use */ diff --git a/frontend/realtime/src/KonvaApp/NodesManagers/ToolbarNodeUtilities.ts b/frontend/realtime/src/KonvaApp/NodesManagers/ToolbarNodeUtilities.ts index 8e09c585be..0c14479ac3 100644 --- a/frontend/realtime/src/KonvaApp/NodesManagers/ToolbarNodeUtilities.ts +++ b/frontend/realtime/src/KonvaApp/NodesManagers/ToolbarNodeUtilities.ts @@ -94,7 +94,8 @@ export function getTextNodeButtonStates( buttonName === ButtonNames.AI_STYLIZE || buttonName === ButtonNames.SEGMENTATION || buttonName === ButtonNames.DOWNLOAD || - buttonName === ButtonNames.CHROMA + buttonName === ButtonNames.CHROMA || + buttonName === ButtonNames.REMOVE_BACKGROUND ) { buttonStates[buttonName] = { disabled: true, @@ -132,7 +133,8 @@ export function getVideoNodeButtonStates( } // soon to come feature is disabled if (buttonName === ButtonNames.AI_STYLIZE || - buttonName === ButtonNames.COLOR + buttonName === ButtonNames.COLOR || + buttonName === ButtonNames.REMOVE_BACKGROUND ) { buttonStates[buttonName] = { disabled: true, @@ -174,7 +176,8 @@ export function getShapeNodeButtonStates( buttonName === ButtonNames.AI_STYLIZE || buttonName === ButtonNames.SEGMENTATION || buttonName === ButtonNames.DOWNLOAD || - buttonName === ButtonNames.CHROMA + buttonName === ButtonNames.CHROMA || + buttonName === ButtonNames.REMOVE_BACKGROUND ) { buttonStates[buttonName] = { disabled: true, @@ -218,7 +221,8 @@ export function getMultiSelectButtonStates( buttonName === ButtonNames.SEGMENTATION || buttonName === ButtonNames.DOWNLOAD || buttonName === ButtonNames.CHROMA || - buttonName === ButtonNames.COLOR + buttonName === ButtonNames.COLOR || + buttonName === ButtonNames.REMOVE_BACKGROUND ) { buttonStates[buttonName] = { disabled: true, diff --git a/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts b/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts index a1f151596c..4d681ec479 100644 --- a/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts +++ b/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts @@ -25,6 +25,7 @@ import { isLoadingVisible, loadingProgress, } from "~/signals/uiEvents/loadingIndicator"; +import {DecodeBase64ToImage} from "~/utilities/DecodeBase64ToImage.ts"; interface ServerSettings { model_path: string; @@ -285,7 +286,7 @@ export class RealTimeDrawEngine { this.client.onResult(async (response) => { try { // Convert the base64 image to ImageBitmap and update the preview - const bitmap = await this.base64ToImageBitmap(response.image); + const bitmap = await DecodeBase64ToImage(response.image); this.outputBitmap = bitmap; this.previewCanvas.image(bitmap); this.mediaLayerRef.batchDraw(); @@ -895,35 +896,6 @@ export class RealTimeDrawEngine { return base64String.split(",")[1]; } - private async base64ToImageBitmap( - base64String: string, - ): Promise { - // 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; - }); - } - public async saveOutput() { if (!this.outputBitmap) { console.error("No output bitmap available to save"); @@ -1007,7 +979,7 @@ export class RealTimeDrawEngine { }); //console.log(base64BitmapResponse); - const decoded = await this.base64ToImageBitmap( + const decoded = await DecodeBase64ToImage( base64BitmapResponse as string, ); diff --git a/frontend/realtime/src/components/features/ToolbarNode/data.ts b/frontend/realtime/src/components/features/ToolbarNode/data.ts index c215e5ed8b..f8bc6a8d71 100644 --- a/frontend/realtime/src/components/features/ToolbarNode/data.ts +++ b/frontend/realtime/src/components/features/ToolbarNode/data.ts @@ -22,6 +22,11 @@ export const ToolbarNodeButtonData = [ icon: faPalette, tooltip: "Color", }, + { + name: ButtonNames.REMOVE_BACKGROUND, + icon: faScalpelLineDashed, + tooltip: "Remove Background", + }, { name: ButtonNames.CHROMA, icon: faTransporter, diff --git a/frontend/realtime/src/components/features/ToolbarNode/enums.ts b/frontend/realtime/src/components/features/ToolbarNode/enums.ts index e4c559695a..7401d94e92 100644 --- a/frontend/realtime/src/components/features/ToolbarNode/enums.ts +++ b/frontend/realtime/src/components/features/ToolbarNode/enums.ts @@ -8,4 +8,5 @@ export enum ToolbarNodeButtonNames { MOVE_LAYER_DOWN = "MOVE_LAYER_DOWN", DOWNLOAD = "DOWNLOAD", COLOR = "COLOR", + REMOVE_BACKGROUND = "REMOVE_BACKGROUND", } diff --git a/frontend/realtime/src/utilities/DecodeBase64ToImage.ts b/frontend/realtime/src/utilities/DecodeBase64ToImage.ts new file mode 100644 index 0000000000..b568808afd --- /dev/null +++ b/frontend/realtime/src/utilities/DecodeBase64ToImage.ts @@ -0,0 +1,30 @@ +/** + * 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 => { + // 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; + }); +} diff --git a/frontend/realtime/src/utilities/EncodeImageBitmapToBase64.ts b/frontend/realtime/src/utilities/EncodeImageBitmapToBase64.ts new file mode 100644 index 0000000000..31f8ed79c3 --- /dev/null +++ b/frontend/realtime/src/utilities/EncodeImageBitmapToBase64.ts @@ -0,0 +1,22 @@ + +export const EncodeImageBitmapToBase64 = async (imageBitmap: ImageBitmap): Promise => { + // Create a temporary canvas + const canvas = document.createElement("canvas"); + canvas.width = imageBitmap.width; + canvas.height = imageBitmap.height; + + // Draw the ImageBitmap onto the canvas + const ctx = canvas.getContext("2d"); + if (!ctx) throw new Error("Failed to get 2D context"); + + ctx.drawImage(imageBitmap, 0, 0); + + // Convert to base64 + const base64String = canvas.toDataURL("image/png"); + + // Clean up + canvas.remove(); + + // Remove the data:image/png;base64, prefix if you want just the base64 string + return base64String.split(",")[1]; +} From f3babbcb19f6104fec76de8dce9ddfb4b14f9854 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 20 Mar 2025 10:53:34 -0400 Subject: [PATCH 2/6] if statements -> switch statements --- crates/desktop/src-tauri/tauri.conf.json | 5 +- .../NodesManagers/ToolbarNodeUtilities.ts | 213 +++++++++--------- 2 files changed, 107 insertions(+), 111 deletions(-) diff --git a/crates/desktop/src-tauri/tauri.conf.json b/crates/desktop/src-tauri/tauri.conf.json index efc6948567..f408df7602 100644 --- a/crates/desktop/src-tauri/tauri.conf.json +++ b/crates/desktop/src-tauri/tauri.conf.json @@ -32,9 +32,6 @@ "icons/128x128@2x.png", "icons/icon.icns", "icons/icon.ico" - ], - "resources": { - "../../../windows_dlls/*.dll": "." - } + ] } } diff --git a/frontend/realtime/src/KonvaApp/NodesManagers/ToolbarNodeUtilities.ts b/frontend/realtime/src/KonvaApp/NodesManagers/ToolbarNodeUtilities.ts index 0c14479ac3..95c55cd7c5 100644 --- a/frontend/realtime/src/KonvaApp/NodesManagers/ToolbarNodeUtilities.ts +++ b/frontend/realtime/src/KonvaApp/NodesManagers/ToolbarNodeUtilities.ts @@ -46,29 +46,29 @@ export function getImageNodeButtonStates( return buttonStates; } - // hidden buttons - if ( - buttonName === ButtonNames.AI_STYLIZE || - buttonName === ButtonNames.SEGMENTATION || - buttonName === ButtonNames.DOWNLOAD || - buttonName === ButtonNames.CHROMA || - buttonName === ButtonNames.COLOR - ) { - buttonStates[buttonName] = { - disabled: true, - hidden: true, - active: false, - }; - return buttonStates; - } + switch (buttonName) { + case ButtonNames.AI_STYLIZE: + case ButtonNames.SEGMENTATION: + case ButtonNames.DOWNLOAD: + case ButtonNames.CHROMA: + case ButtonNames.COLOR: + // hidden buttons + buttonStates[buttonName] = { + disabled: true, + hidden: true, + active: false, + }; + return buttonStates; - // all other buttons - buttonStates[buttonName] = { - disabled: false, - hidden: false, - active: false, - }; - return buttonStates; + default: + // all other buttons + buttonStates[buttonName] = { + disabled: false, + hidden: false, + active: false, + }; + return buttonStates; + } }, {} as ContextualToolbarProps["buttonStates"], ); @@ -89,29 +89,29 @@ export function getTextNodeButtonStates( return buttonStates; } - // hidden buttons - if ( - buttonName === ButtonNames.AI_STYLIZE || - buttonName === ButtonNames.SEGMENTATION || - buttonName === ButtonNames.DOWNLOAD || - buttonName === ButtonNames.CHROMA || - buttonName === ButtonNames.REMOVE_BACKGROUND - ) { - buttonStates[buttonName] = { - disabled: true, - hidden: true, - active: false, - }; - return buttonStates; - } + switch (buttonName) { + case ButtonNames.AI_STYLIZE: + case ButtonNames.SEGMENTATION: + case ButtonNames.DOWNLOAD: + case ButtonNames.CHROMA: + case ButtonNames.REMOVE_BACKGROUND: + // hidden buttons + buttonStates[buttonName] = { + disabled: true, + hidden: true, + active: false, + }; + return buttonStates; - // all other buttons - buttonStates[buttonName] = { - disabled: false, - hidden: false, - active: false, - }; - return buttonStates; + default: + // all other buttons + buttonStates[buttonName] = { + disabled: false, + hidden: false, + active: false, + }; + return buttonStates; + } }, {} as ContextualToolbarProps["buttonStates"], ); @@ -131,25 +131,26 @@ export function getVideoNodeButtonStates( }; return buttonStates; } - // soon to come feature is disabled - if (buttonName === ButtonNames.AI_STYLIZE || - buttonName === ButtonNames.COLOR || - buttonName === ButtonNames.REMOVE_BACKGROUND - ) { - buttonStates[buttonName] = { - disabled: true, - hidden: true, - active: false, - }; - return buttonStates; + switch (buttonName) { + case ButtonNames.AI_STYLIZE: + case ButtonNames.COLOR: + case ButtonNames.REMOVE_BACKGROUND: + // soon to come feature is disabled + buttonStates[buttonName] = { + disabled: true, + hidden: true, + active: false, + }; + return buttonStates; + default: + // all other buttons + buttonStates[buttonName] = { + disabled: false, + hidden: false, + active: false, + }; + return buttonStates; } - // all other buttons - buttonStates[buttonName] = { - disabled: false, - hidden: false, - active: false, - }; - return buttonStates; }, {} as ContextualToolbarProps["buttonStates"], ); @@ -171,29 +172,28 @@ export function getShapeNodeButtonStates( return buttonStates; } - // hidden buttons - if ( - buttonName === ButtonNames.AI_STYLIZE || - buttonName === ButtonNames.SEGMENTATION || - buttonName === ButtonNames.DOWNLOAD || - buttonName === ButtonNames.CHROMA || - buttonName === ButtonNames.REMOVE_BACKGROUND - ) { - buttonStates[buttonName] = { - disabled: true, - hidden: true, - active: false, - }; - return buttonStates; + switch (buttonName) { + case ButtonNames.AI_STYLIZE: + case ButtonNames.SEGMENTATION: + case ButtonNames.DOWNLOAD: + case ButtonNames.CHROMA: + case ButtonNames.REMOVE_BACKGROUND: + // hidden buttons + buttonStates[buttonName] = { + disabled: true, + hidden: true, + active: false, + }; + return buttonStates; + default: + // all other buttons + buttonStates[buttonName] = { + disabled: false, + hidden: false, + active: false, + }; + return buttonStates; } - - // all other buttons - buttonStates[buttonName] = { - disabled: false, - hidden: false, - active: false, - }; - return buttonStates; }, {} as ContextualToolbarProps["buttonStates"], ); @@ -215,30 +215,29 @@ export function getMultiSelectButtonStates( return buttonStates; } - // hidden buttons - if ( - buttonName === ButtonNames.AI_STYLIZE || - buttonName === ButtonNames.SEGMENTATION || - buttonName === ButtonNames.DOWNLOAD || - buttonName === ButtonNames.CHROMA || - buttonName === ButtonNames.COLOR || - buttonName === ButtonNames.REMOVE_BACKGROUND - ) { - buttonStates[buttonName] = { - disabled: true, - hidden: true, - active: false, - }; - return buttonStates; + switch (buttonName) { + case ButtonNames.AI_STYLIZE: + case ButtonNames.SEGMENTATION: + case ButtonNames.DOWNLOAD: + case ButtonNames.CHROMA: + case ButtonNames.COLOR: + case ButtonNames.REMOVE_BACKGROUND: + // hidden buttons + buttonStates[buttonName] = { + disabled: true, + hidden: true, + active: false, + }; + return buttonStates; + default: + // all other buttons + buttonStates[buttonName] = { + disabled: false, + hidden: false, + active: false, + }; + return buttonStates; } - - // all other buttons - buttonStates[buttonName] = { - disabled: false, - hidden: false, - active: false, - }; - return buttonStates; }, {} as ContextualToolbarProps["buttonStates"], ); From d11945b6854ae3f91923b7756e8f6ecd1881ddf2 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 20 Mar 2025 11:07:34 -0400 Subject: [PATCH 3/6] download DIS model --- .../src-tauri/src/endpoints/remove_background_endpoint.rs | 3 ++- crates/desktop/src-tauri/src/events/notification_event.rs | 5 +++-- .../src/ml/weights_registry/weight_descriptor.rs | 2 ++ .../desktop/src-tauri/src/ml/weights_registry/weights.rs | 8 ++++++++ crates/desktop/src-tauri/src/threads/downloader_thread.rs | 4 +++- 5 files changed, 18 insertions(+), 4 deletions(-) diff --git a/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs index 487315b866..e487519b34 100644 --- a/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs @@ -10,6 +10,7 @@ use log::{error, info}; use ndarray::{Array, CowArray}; use ort::Value; use tauri::{AppHandle, State}; +use crate::ml::weights_registry::weights::{DIS_MEDIUM_ONNX, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; /// This handler removes the background from an image. #[tauri::command] @@ -28,7 +29,7 @@ pub async fn remove_background( // sudo apt-get install libcudnn8-dev libcudnn8 // update-alternatives: warning: forcing reinstallation of alternative /usr/include/x86_64-linux-gnu/cudnn_v9.h because link group libcudnn is broken // update-alternatives: using /usr/include/x86_64-linux-gnu/cudnn_v8.h to provide /usr/include/cudnn.h (libcudnn) in manual mode - let model_path = app_data_root.weights_dir().path().join("medium.onnx"); // TODO + let model_path = app_data_root.weights_dir().weight_path(&DIS_MEDIUM_ONNX); let session = onnx_session(model_path.to_str().unwrap()) // TODO .map_err(|err| format!("failure to start onnx session: {:?}", err))?; diff --git a/crates/desktop/src-tauri/src/events/notification_event.rs b/crates/desktop/src-tauri/src/events/notification_event.rs index 747d68ea4c..5bfaa5c491 100644 --- a/crates/desktop/src-tauri/src/events/notification_event.rs +++ b/crates/desktop/src-tauri/src/events/notification_event.rs @@ -4,10 +4,11 @@ use serde::Serialize; #[derive(Clone, Copy, Serialize)] #[serde(rename_all = "snake_case")] pub enum NotificationModelType { - Unet, - Vae, + ImageSegmentation, TextEncoder, TextTokenizer, + Unet, + Vae, } #[derive(Clone, Serialize)] diff --git a/crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor.rs b/crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor.rs index 4d4dbfcf4c..dca41c0099 100644 --- a/crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor.rs +++ b/crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor.rs @@ -5,6 +5,7 @@ #[derive(Clone, Copy)] pub enum WeightFunction { + ImageSegmentation, TextEncoder, TextTokenizer, Unet, @@ -28,6 +29,7 @@ pub struct WeightDescriptor { impl WeightFunction { pub fn as_descriptive_str(&self) -> &'static str { match self { + WeightFunction::ImageSegmentation => "image segmentation", WeightFunction::TextEncoder => "text encoder", WeightFunction::TextTokenizer => "text tokenizer", WeightFunction::Unet => "UNET", diff --git a/crates/desktop/src-tauri/src/ml/weights_registry/weights.rs b/crates/desktop/src-tauri/src/ml/weights_registry/weights.rs index 4e44e453fa..f2d35f50da 100644 --- a/crates/desktop/src-tauri/src/ml/weights_registry/weights.rs +++ b/crates/desktop/src-tauri/src/ml/weights_registry/weights.rs @@ -28,6 +28,14 @@ pub const SIMIANLUO_LCM_DREAMSHAPER_V7_UNET : WeightDescriptor = weight!( WeightFunction::Unet ); +/// Image segmentation and background removal +/// dis_medium.onnx +pub const DIS_MEDIUM_ONNX : WeightDescriptor = weight!( + "Dichotomous Image Segmentation Medium", + "dis_medium.onnx", + WeightFunction::ImageSegmentation +); + #[deprecated] pub const SDXL_TURBO_CLIP_TEXT_ENCODER : WeightDescriptor = weight!( "SDXL Turbo CLIP", diff --git a/crates/desktop/src-tauri/src/threads/downloader_thread.rs b/crates/desktop/src-tauri/src/threads/downloader_thread.rs index 3a9992fc7d..8364ff02ec 100644 --- a/crates/desktop/src-tauri/src/threads/downloader_thread.rs +++ b/crates/desktop/src-tauri/src/threads/downloader_thread.rs @@ -12,7 +12,7 @@ use tauri::{AppHandle, Emitter}; use tempfile::NamedTempFile; use tokio::task::JoinHandle; use crate::ml::weights_registry::weight_descriptor::WeightFunction; -use crate::ml::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2, SDXL_TURBO_UNET, SDXL_TURBO_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; +use crate::ml::weights_registry::weights::{CLIP_JSON, DIS_MEDIUM_ONNX, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2, SDXL_TURBO_UNET, SDXL_TURBO_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; const NOTIFICATION_CHANNEL_NAME : &str = "notification"; const MAX_FILES : usize = 8; @@ -45,6 +45,7 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> ! download_queue.push_back(LYKON_DEAMSHAPER_7_VAE); download_queue.push_back(SIMIANLUO_LCM_DREAMSHAPER_V7_UNET); download_queue.push_back(SDXL_TURBO_CLIP_TEXT_ENCODER); // TODO(bt): Why is this still needed? + download_queue.push_back(DIS_MEDIUM_ONNX); //download_queue.push_back(SDXL_TURBO_CLIP_TEXT_ENCODER_2); // TODO(bt): REMOVE //download_queue.push_back(SDXL_TURBO_VAE); // TODO(bt): REMOVE //download_queue.push_back(SDXL_TURBO_UNET); // TODO(bt): REMOVE @@ -80,6 +81,7 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> ! WeightFunction::TextTokenizer => NotificationModelType::TextTokenizer, WeightFunction::Unet => NotificationModelType::Unet, WeightFunction::Vae => NotificationModelType::Vae, + WeightFunction::ImageSegmentation => NotificationModelType::ImageSegmentation, }; let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadStarted { From 5fdb4d1024c296077f0825e272de5da0033ef7b1 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 20 Mar 2025 11:59:43 -0400 Subject: [PATCH 4/6] disable msi builds due to 2.0 gb limit for 32bit --- crates/desktop/src-tauri/tauri.conf.json | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/crates/desktop/src-tauri/tauri.conf.json b/crates/desktop/src-tauri/tauri.conf.json index f408df7602..718a24832a 100644 --- a/crates/desktop/src-tauri/tauri.conf.json +++ b/crates/desktop/src-tauri/tauri.conf.json @@ -25,13 +25,16 @@ }, "bundle": { "active": true, - "targets": "all", + "targets": ["deb", "rpm", "nsis"], "icon": [ "icons/32x32.png", "icons/128x128.png", "icons/128x128@2x.png", "icons/icon.icns", "icons/icon.ico" - ] + ], + "resources": { + "../../../windows_dlls/*.dll": "." + } } } From 32cb39a643a1bf388b9eca565b0d474a7272360f Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 20 Mar 2025 23:28:59 -0400 Subject: [PATCH 5/6] text_to_image endpoint --- .../src-tauri/binary_includes/1024.png | Bin 0 -> 8630 bytes crates/desktop/src-tauri/src/endpoints/mod.rs | 3 +- ...endpoint.rs => realtime_image_endpoint.rs} | 0 .../src/endpoints/text_to_image_endpoint.rs | 99 ++++++++++++++++++ crates/desktop/src-tauri/src/lib.rs | 4 +- crates/desktop/src-tauri/tauri.conf.json | 5 +- 6 files changed, 105 insertions(+), 6 deletions(-) create mode 100644 crates/desktop/src-tauri/binary_includes/1024.png rename crates/desktop/src-tauri/src/endpoints/{image_endpoint.rs => realtime_image_endpoint.rs} (100%) create mode 100644 crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs diff --git a/crates/desktop/src-tauri/binary_includes/1024.png b/crates/desktop/src-tauri/binary_includes/1024.png new file mode 100644 index 0000000000000000000000000000000000000000..9f746687dc4322ed7dafcbf34737ed681a7bc018 GIT binary patch literal 8630 zcmeI2ze_?<6vw}&Wtu_-Q3!!iv=kwYA;dGy=kTPYrOjw?73_z!^e_pUszGdMFrq06 zS`u6m8mgt;TB;=zL0ZM3bKm(J!sR@8oO|wj?{fLv?>+BuxJN7TsMqat1MtT5Z~|~p z(}BSrGNo3ssJxVBXu#d^Kxd`en^9%xgSSqHWQKqCn( zoB^Zh^U7+F9=Ou_dIp%ARJSdAwYWo_&aIdcasJpx?0w^nFa4B;L@cbWW#5`LZHS!5 zK=tAK^RlSd%e(LQcKQU!2VIseJy`~(pZk&Xa7X_-7M2xr4;O+lHq0v!m|~mM*;Y^3niDo|EUC&iU5BdBO_vNH3`7UkR*=0Q z@xT&-#6*dS5)&mRb~8tsx%O4)tbuhFfLEfMv+33hE`Wl!{c1Gnj0F&j#KX6tWbO|U CHh4Aw literal 0 HcmV?d00001 diff --git a/crates/desktop/src-tauri/src/endpoints/mod.rs b/crates/desktop/src-tauri/src/endpoints/mod.rs index d365df449b..f49bdce06e 100644 --- a/crates/desktop/src-tauri/src/endpoints/mod.rs +++ b/crates/desktop/src-tauri/src/endpoints/mod.rs @@ -1,5 +1,6 @@ pub mod download_models; pub mod flip_image; -pub mod image_endpoint; +pub mod realtime_image_endpoint; pub mod remove_background_endpoint; pub mod test_counter; +pub mod text_to_image_endpoint; diff --git a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs similarity index 100% rename from crates/desktop/src-tauri/src/endpoints/image_endpoint.rs rename to crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs diff --git a/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs new file mode 100644 index 0000000000..081625fb81 --- /dev/null +++ b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs @@ -0,0 +1,99 @@ +use crate::ml::model_cache::ModelCache; +use crate::ml::prompt_cache::PromptCache; +use crate::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args}; +use crate::state::app_config::AppConfig; +use crate::state::app_dir::AppDataRoot; +use crate::utils::image::decode_base64_image::decode_base64_image; +use crate::utils::image::encode_rgb_image_base64_png::encode_rgb_image_base64_png; +use image::imageops::FilterType; +use image::{DynamicImage, ImageReader, RgbImage}; +use log::{error, info}; +use once_cell::sync::Lazy; +use std::io::Cursor; +use std::path::PathBuf; +use tauri::{AppHandle, State}; + +const RANDOM_SEED: u32 = 42; + +const STRENGTH: f64 = 1.0; +const PNG_BYTES : &[u8] = include_bytes!("../../binary_includes/1024.png"); + +// NB: We're just hacking a temporary "text to image" test endpoint leveraging the existing image to image modality. +static IMAGE : Lazy = Lazy::new(||{ + ImageReader::new(Cursor::new(PNG_BYTES)) + .with_guessed_format() + .expect("failed to detect png format") + .decode() + .expect("failed to decode image") +}); + +#[tauri::command] +pub async fn text_to_image( + prompt: String, + model_config: State<'_, AppConfig>, + model_cache: State<'_, ModelCache>, + prompt_cache: State<'_, PromptCache>, + app_data_root: State<'_, AppDataRoot>, + app: AppHandle, +) -> Result { + info!("text_to_image endpoint called."); + + let image = IMAGE.resize(512, 512, FilterType::CatmullRom); + + let result = do_infer_image( + &prompt, + image, + Some(STRENGTH), + &model_config, + &model_cache, + prompt_cache, + app, + &app_data_root + ).await; + + let image = match result { + Ok(image) => image, + Err(err) => { + error!("There was an error generating the image: {:?}", err); + return Err(format!("There was an error generating the image: {}", err)); + } + }; + + let bytes = encode_rgb_image_base64_png(image) + .map_err(|err| format!("failure to encode image: {:?}", err))?; + + info!("Inference successful; image converted to base64, serving back to browser..."); + + Ok(bytes) +} + +async fn do_infer_image( + prompt: &str, + image: DynamicImage, + strength: Option, + config: &AppConfig, + model_cache: &ModelCache, + prompt_cache: State<'_, PromptCache>, + app: AppHandle, + app_data_root: &AppDataRoot, +) -> Result { + + let args = Args { + image: &image, + prompt: prompt.to_string(), + uncond_prompt: "".to_string(), + model_cache, + configs: config, + prompt_cache: &prompt_cache, + i2i_strength: strength, + cfg_scale: config.cfg_scale, + app: &app, + app_data_root, + use_flash_attn: true, + }; + + let image = lcm_pipeline(args) + .map_err(|err| format!("failure to encode image: {:?}", err))?; + + Ok(image) +} diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index a5f00e13e2..653655c2a0 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/crates/desktop/src-tauri/src/lib.rs @@ -7,9 +7,10 @@ pub mod transfer; pub mod utils; use crate::endpoints::download_models::download_models; -use crate::endpoints::image_endpoint::infer_image; +use crate::endpoints::realtime_image_endpoint::infer_image; use crate::endpoints::remove_background_endpoint::remove_background; use crate::endpoints::test_counter::test_counter; +use crate::endpoints::text_to_image_endpoint::text_to_image; use crate::ml::model_cache::ModelCache; use crate::ml::prompt_cache::PromptCache; use crate::state::app_config::AppConfig; @@ -66,6 +67,7 @@ pub fn run() { infer_image, remove_background, test_counter, + text_to_image, ]) .run(tauri::generate_context!()) .expect("error while running tauri application"); diff --git a/crates/desktop/src-tauri/tauri.conf.json b/crates/desktop/src-tauri/tauri.conf.json index 718a24832a..17c4b70706 100644 --- a/crates/desktop/src-tauri/tauri.conf.json +++ b/crates/desktop/src-tauri/tauri.conf.json @@ -32,9 +32,6 @@ "icons/128x128@2x.png", "icons/icon.icns", "icons/icon.ico" - ], - "resources": { - "../../../windows_dlls/*.dll": "." - } + ] } } From 08590d0c61172a07d1133b2a7b159284337f74ba Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 21 Mar 2025 01:11:34 -0400 Subject: [PATCH 6/6] log environment details --- Cargo.lock | 80 ++++++++++++++++--- crates/desktop/src-tauri/Cargo.toml | 8 +- crates/desktop/src-tauri/src/lib.rs | 4 +- .../src/utils/log_environment_details.rs | 11 +++ crates/desktop/src-tauri/src/utils/mod.rs | 3 +- 5 files changed, 90 insertions(+), 16 deletions(-) create mode 100644 crates/desktop/src-tauri/src/utils/log_environment_details.rs diff --git a/Cargo.lock b/Cargo.lock index 7dc762487b..4c877cc071 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -677,6 +677,7 @@ dependencies = [ "log", "memory_caching", "ndarray", + "nvidia-checker", "once_cell", "ort", "primitives", @@ -1883,6 +1884,18 @@ dependencies = [ "toml 0.8.14", ] +[[package]] +name = "carrot_utils" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f34277d609ff503187aa8e8604913b103e2eb8f893222a4948748f22527252e" +dependencies = [ + "chrono", + "colored", + "log", + "simplelog", +] + [[package]] name = "cast" version = "0.3.0" @@ -2060,9 +2073,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.4.8" +version = "4.5.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2275f18819641850fa26c89acc84d465c1bf91ce57bc2748b28c420473352f64" +checksum = "0fbb260a053428790f3de475e304ff84cdbc4face759ea7a3e64c1edd938a7fc" dependencies = [ "clap_builder", "clap_derive", @@ -2070,23 +2083,23 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.4.8" +version = "4.5.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07cdf1b148b25c1e1f7a42225e30a0d99a615cd4637eae7365548dd4529b95bc" +checksum = "64b17d7ea74e9f833c7dbf2cbe4fb12ff26783eda4782a8975b72f895c9b4d99" dependencies = [ "anstream", "anstyle", "clap_lex", - "strsim 0.10.0", + "strsim 0.11.1", ] [[package]] name = "clap_derive" -version = "4.4.7" +version = "4.5.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf9804afaaf59a91e75b022a30fb7229a7901f60c755489cc61c9b423b836442" +checksum = "501d359d5f3dcaf6ecdeee48833ae73ec6e42723a1e52419c79abf9507eec0a0" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.87", @@ -2094,9 +2107,9 @@ dependencies = [ [[package]] name = "clap_lex" -version = "0.6.0" +version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "702fc72eb24e5a1e48ce58027a675bc24edd52096d5397d4aea7c6dd9eca0bd1" +checksum = "f46ad14479a25103f283c0f10005961cf086d8dc42205bb44c46ac563475dca6" [[package]] name = "cloud_storage" @@ -2189,6 +2202,16 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b6a852b24ab71dffc585bcb46eaf7959d175cb865a7152e35b348d1b2960422" +[[package]] +name = "colored" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" +dependencies = [ + "lazy_static", + "windows-sys 0.59.0", +] + [[package]] name = "combine" version = "3.8.1" @@ -7335,6 +7358,23 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" +[[package]] +name = "nvidia-checker" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8b584e4b0be14d4e3f8b779d79477aa0f36658d7a535c42abc725013d0c8951" +dependencies = [ + "carrot_utils", + "chrono", + "clap", + "colored", + "dirs 5.0.1", + "log", + "serde", + "serde_derive", + "toml 0.8.14", +] + [[package]] name = "objc" version = "0.2.7" @@ -10450,6 +10490,17 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" +[[package]] +name = "simplelog" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16257adbfaef1ee58b1363bdc0664c9b8e1e30aed86049635fb5f147d065a9c0" +dependencies = [ + "log", + "termcolor", + "time", +] + [[package]] name = "siphasher" version = "0.3.11" @@ -11980,6 +12031,15 @@ dependencies = [ "safetensors", ] +[[package]] +name = "termcolor" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755" +dependencies = [ + "winapi-util", +] + [[package]] name = "termtree" version = "0.4.1" diff --git a/crates/desktop/src-tauri/Cargo.toml b/crates/desktop/src-tauri/Cargo.toml index 5da606b9e3..a3a354df2e 100644 --- a/crates/desktop/src-tauri/Cargo.toml +++ b/crates/desktop/src-tauri/Cargo.toml @@ -39,10 +39,14 @@ const_format = "0.2.34" directories = "6.0.0" dirs = "6.0.0" futures = "0.3" +half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] } hf-hub = "0.3.2" image = "0.25.5" log = "0.4" +ndarray = "0.15.6" +nvidia-checker = "0.1.2" once_cell = "1.20.3" +ort = { version = "1.16.2", features = ["cuda"] } rand = "0.8.5" # NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls reqwest = { version = "0.12.12", features = ["stream"], default-features = false } @@ -56,11 +60,7 @@ tempfile = "3.18.0" tokenizers = "0.15.2" tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] } tokio-util = { version = "0.7", features = ["codec"] } -half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] } -# NB: From `rust-background-removal` -ndarray = "0.15.6" -ort = { version = "1.16.2", features = ["cuda"] } # tracing-subscriber = { version = "0.3", features = [ "env-filter", "fmt" ] } [features] diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index 653655c2a0..e35b9cb46a 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/crates/desktop/src-tauri/src/lib.rs @@ -15,6 +15,7 @@ use crate::ml::model_cache::ModelCache; use crate::ml::prompt_cache::PromptCache; use crate::state::app_config::AppConfig; use crate::threads::downloader_thread::downloader_thread; +use crate::utils::log_environment_details::log_environment_details; use tauri_plugin_log::Target; use tauri_plugin_log::TargetKind; @@ -45,6 +46,7 @@ pub fn run() { .targets(vec![Target::new(TargetKind::Stdout)]) .build()) .setup(|app| { + log_environment_details(); //if cfg!(debug_assertions) { // app.handle().plugin( // tauri_plugin_log::Builder::default() @@ -66,7 +68,7 @@ pub fn run() { download_models, infer_image, remove_background, - test_counter, + test_counter, text_to_image, ]) .run(tauri::generate_context!()) diff --git a/crates/desktop/src-tauri/src/utils/log_environment_details.rs b/crates/desktop/src-tauri/src/utils/log_environment_details.rs new file mode 100644 index 0000000000..3b54104246 --- /dev/null +++ b/crates/desktop/src-tauri/src/utils/log_environment_details.rs @@ -0,0 +1,11 @@ +use log::info; +use nvidia_checker::{get_cuda_version, get_cudnn_version, get_kernel_version, get_nvidia_driver_version, get_os_version, get_tensorrt_version}; + +pub fn log_environment_details() { + info!("NVIDIA driver version: {}", get_nvidia_driver_version(false)); + info!("CUDA version: {}", get_cuda_version(false)); + info!("cuDNN version: {}", get_cudnn_version(false)); + info!("TensorRT version: {}", get_tensorrt_version(false)); + info!("OS version: {}", get_os_version(false)); + info!("Kernel version: {}", get_kernel_version(false)); +} diff --git a/crates/desktop/src-tauri/src/utils/mod.rs b/crates/desktop/src-tauri/src/utils/mod.rs index b1b6b4f003..3e475ef135 100644 --- a/crates/desktop/src-tauri/src/utils/mod.rs +++ b/crates/desktop/src-tauri/src/utils/mod.rs @@ -1 +1,2 @@ -pub mod image; \ No newline at end of file +pub mod image; +pub mod log_environment_details; \ No newline at end of file