working on file browser

2024-08-22 22:28:29 -07:00
parent 852542d7df
commit 3bb746730d
3 changed files with 369 additions and 0 deletions
@@ -57,6 +57,7 @@ from .nodes.FL_SchedulerStrings import FL_SchedulerStrings
 from .nodes.FL_ImageCaptionLayoutPDF import FL_ImageCaptionLayoutPDF
 from .nodes.FL_Dalle3 import FL_Dalle3
 from .nodes.FL_SaveImages import FL_SaveImages
+from .nodes.FL_FileBrowser import FL_FileBrowser



@@ -121,6 +122,7 @@ NODE_CLASS_MAPPINGS = {
    "FL_ImageCaptionLayoutPDF": FL_ImageCaptionLayoutPDF,
    "FL_Dalle3": FL_Dalle3,
    "FL_SaveImages": FL_SaveImages,
+    "FL_FileBrowser": FL_FileBrowser,

 }

@@ -184,6 +186,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
    "FL_ImageCaptionLayoutPDF": "FL Image Caption Layout PDF",
    "FL_Dalle3": "FL Dalle 3",
    "FL_SaveImages": "FL Save Images",
+    "FL_FileBrowser": "FL File Browser",

 }

@@ -0,0 +1,103 @@
+import os
+import json
+from PIL import Image
+import torch
+import numpy as np
+from server import PromptServer
+from aiohttp import web
+import io
+
+
+class FL_FileBrowser:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "root_directory": ("STRING", {"default": "./"}),
+            },
+            "optional": {
+                "selected_file": ("STRING", {"default": ""}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE", "STRING")
+    FUNCTION = "browse_files"
+    OUTPUT_NODE = True
+    CATEGORY = "🏵️Fill Nodes/Image"
+
+    def browse_files(self, root_directory, selected_file=""):
+        if selected_file and os.path.isfile(selected_file):
+            image = Image.open(selected_file)
+            image_tensor = torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+            return (image_tensor, selected_file)
+        else:
+            return (torch.zeros(1, 1, 1, 3), "")
+
+    @classmethod
+    def IS_CHANGED(cls, root_directory, selected_file=""):
+        return selected_file
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, root_directory, selected_file=""):
+        if not os.path.isdir(root_directory):
+            return "Root directory does not exist"
+        if selected_file and not os.path.isfile(selected_file):
+            return "Selected file does not exist"
+        return True
+
+
+def get_directory_structure(path):
+    structure = {"name": os.path.basename(path), "children": [], "path": path}
+    try:
+        with os.scandir(path) as entries:
+            for entry in entries:
+                if entry.is_dir():
+                    structure["children"].append(get_directory_structure(entry.path))
+    except PermissionError:
+        pass
+    return structure
+
+
+def get_file_list(path):
+    return [f for f in os.listdir(path) if
+            os.path.isfile(os.path.join(path, f)) and f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp'))]
+
+
+@PromptServer.instance.routes.post("/fl_file_browser/get_directory_structure")
+async def api_get_directory_structure(request):
+    data = await request.json()
+    path = data.get("path", "./")
+
+    if not os.path.isabs(path):
+        path = os.path.abspath(path)
+
+    if not os.path.exists(path):
+        return web.json_response({"error": "Directory does not exist"}, status=400)
+
+    structure = get_directory_structure(path)
+    files = get_file_list(path)
+    return web.json_response({"structure": structure, "files": files})
+
+
+@PromptServer.instance.routes.post("/fl_file_browser/get_thumbnail")
+async def api_get_thumbnail(request):
+    data = await request.json()
+    path = data.get("path", "./")
+    file = data.get("file", "")
+
+    if not os.path.isabs(path):
+        path = os.path.abspath(path)
+
+    full_path = os.path.join(path, file)
+    if not os.path.exists(full_path):
+        return web.json_response({"error": "File does not exist"}, status=400)
+
+    try:
+        with Image.open(full_path) as img:
+            img.thumbnail((80, 80))
+            buf = io.BytesIO()
+            img.save(buf, format='PNG')
+            buf.seek(0)
+            return web.Response(body=buf.read(), content_type='image/png')
+    except Exception as e:
+        return web.json_response({"error": str(e)}, status=500)
@@ -0,0 +1,263 @@
+import { app } from "../../../scripts/app.js";
+
+app.registerExtension({
+    name: "Comfy.FL_FileBrowser",
+    async nodeCreated(node) {
+        if (node.comfyClass === "FL_FileBrowser") {
+            addFileBrowserUI(node);
+        }
+    }
+});
+
+function addFileBrowserUI(node) {
+    const rootDirectoryWidget = node.widgets.find(w => w.name === "root_directory");
+    const selectedFileWidget = node.widgets.find(w => w.name === "selected_file");
+
+    rootDirectoryWidget.hidden = false;
+    selectedFileWidget.hidden = true;
+
+    const MIN_WIDTH = 400;
+    const MIN_HEIGHT = 450;
+    const TOP_PADDING = 140;
+    const BOTTOM_PADDING = 20;
+    const FOLDER_HEIGHT = 20;
+    const INDENT_WIDTH = 15;
+    const TOP_BAR_HEIGHT = 30;
+    const THUMBNAIL_SIZE = 80;
+    const THUMBNAIL_PADDING = 5;
+
+    let currentDirectory = rootDirectoryWidget.value;
+    let selectedFile = selectedFileWidget.value;
+    let directoryStructure = { name: "root", children: [] };
+    let fileList = [];
+    let thumbnails = {};
+    let scrollOffsetLeft = 0;
+    let scrollOffsetRight = 0;
+
+    async function updateDirectoryStructure() {
+        try {
+            const response = await fetch('/fl_file_browser/get_directory_structure', {
+                method: 'POST',
+                headers: { 'Content-Type': 'application/json' },
+                body: JSON.stringify({ path: currentDirectory })
+            });
+            const data = await response.json();
+            directoryStructure = data.structure;
+            fileList = data.files;
+            loadThumbnails();
+            node.setDirtyCanvas(true);
+        } catch (error) {
+            console.error("Error updating directory structure:", error);
+        }
+    }
+
+    async function loadThumbnails() {
+        thumbnails = {};
+        for (const file of fileList) {
+            const response = await fetch('/fl_file_browser/get_thumbnail', {
+                method: 'POST',
+                headers: { 'Content-Type': 'application/json' },
+                body: JSON.stringify({ path: currentDirectory, file: file })
+            });
+            const blob = await response.blob();
+            thumbnails[file] = await createImageBitmap(blob);
+        }
+        node.setDirtyCanvas(true);
+    }
+
+    function updateSelectedFile(file) {
+        selectedFile = file;
+        selectedFileWidget.value = currentDirectory + '/' + file;
+        node.setDirtyCanvas(true);
+    }
+
+    function goUpDirectory() {
+        const parentDir = currentDirectory.split(/(\\|\/)/g).slice(0, -2).join('');
+        if (parentDir) {
+            currentDirectory = parentDir;
+            rootDirectoryWidget.value = currentDirectory;
+            updateDirectoryStructure();
+        }
+    }
+
+    const refreshButton = node.addWidget("button", "Refresh", null, () => {
+        currentDirectory = rootDirectoryWidget.value;
+        updateDirectoryStructure();
+    });
+
+    rootDirectoryWidget.callback = () => {
+        currentDirectory = rootDirectoryWidget.value;
+        updateDirectoryStructure();
+    };
+
+    node.onDrawBackground = function(ctx) {
+        if (!this.flags.collapsed) {
+            const pos = TOP_PADDING - TOP_BAR_HEIGHT;
+            ctx.fillStyle = "#2A2A2A";
+            ctx.fillRect(0, pos, this.size[0], this.size[1] - pos);
+
+            // Draw top bar
+            ctx.fillStyle = "#3A3A3A";
+            ctx.fillRect(0, pos, this.size[0], TOP_BAR_HEIGHT);
+
+            // Draw back button
+            ctx.fillStyle = "#4A4A4A";
+            ctx.fillRect(5, pos + 5, 60, TOP_BAR_HEIGHT - 10);
+            ctx.fillStyle = "#FFFFFF";
+            ctx.font = "12px Arial";
+            ctx.fillText("← Back", 15, pos + 20);
+
+            // Draw current directory
+            ctx.fillStyle = "#FFFFFF";
+            ctx.font = "12px Arial";
+            ctx.fillText(currentDirectory, 75, pos + 20);
+
+            const midX = this.size[0] / 2;
+
+            // Set up clipping regions for scrolling
+            ctx.save();
+            ctx.beginPath();
+            ctx.rect(0, TOP_PADDING, midX, this.size[1] - TOP_PADDING - BOTTOM_PADDING);
+            ctx.clip();
+            drawDirectoryStructure(ctx, 10, TOP_PADDING - scrollOffsetLeft, directoryStructure);
+            ctx.restore();
+
+            ctx.save();
+            ctx.beginPath();
+            ctx.rect(midX, TOP_PADDING, this.size[0] - midX, this.size[1] - TOP_PADDING - BOTTOM_PADDING);
+            ctx.clip();
+            drawThumbnails(ctx, midX, TOP_PADDING - scrollOffsetRight, this.size[0] / 2 - 10, this.size[1] - TOP_PADDING - BOTTOM_PADDING);
+            ctx.restore();
+        }
+    };
+
+    function drawDirectoryStructure(ctx, x, y, structure, level = 0) {
+        const folderIcon = "📁";
+        ctx.font = "14px Arial";
+        ctx.fillStyle = "#ffffff";
+
+        for (const item of structure.children) {
+            const xPos = x + INDENT_WIDTH * level;
+            ctx.fillText(`${folderIcon} ${item.name}`, xPos, y);
+            y += FOLDER_HEIGHT;
+
+            if (item.children) {
+                y = drawDirectoryStructure(ctx, x, y, item, level + 1);
+            }
+        }
+
+        return y;
+    }
+
+    function drawThumbnails(ctx, x, y, width, height) {
+        ctx.fillStyle = "#3A3A3A";
+        ctx.fillRect(x, y, width, height);
+
+        const thumbnailsPerRow = Math.floor(width / (THUMBNAIL_SIZE + THUMBNAIL_PADDING));
+        fileList.forEach((file, index) => {
+            const row = Math.floor(index / thumbnailsPerRow);
+            const col = index % thumbnailsPerRow;
+            const xPos = x + col * (THUMBNAIL_SIZE + THUMBNAIL_PADDING) + THUMBNAIL_PADDING;
+            const yPos = y + row * (THUMBNAIL_SIZE + THUMBNAIL_PADDING) + THUMBNAIL_PADDING;
+
+            if (thumbnails[file]) {
+                ctx.drawImage(thumbnails[file], xPos, yPos, THUMBNAIL_SIZE, THUMBNAIL_SIZE);
+            } else {
+                ctx.fillStyle = "#4A4A4A";
+                ctx.fillRect(xPos, yPos, THUMBNAIL_SIZE, THUMBNAIL_SIZE);
+            }
+
+            if (file === selectedFile) {
+                ctx.strokeStyle = "#4a90e2";
+                ctx.lineWidth = 2;
+                ctx.strokeRect(xPos, yPos, THUMBNAIL_SIZE, THUMBNAIL_SIZE);
+            }
+
+            ctx.fillStyle = "#FFFFFF";
+            ctx.font = "10px Arial";
+            ctx.fillText(file.substring(0, 10) + (file.length > 10 ? "..." : ""), xPos, yPos + THUMBNAIL_SIZE + 12);
+        });
+    }
+
+    node.onMouseDown = function(event) {
+        const pos = TOP_PADDING - TOP_BAR_HEIGHT;
+        const localY = event.canvasY - this.pos[1] - pos;
+        const localX = event.canvasX - this.pos[0];
+
+        if (localY >= 0 && localY <= TOP_BAR_HEIGHT) {
+            // Click on top bar
+            if (localX >= 5 && localX <= 65 && localY >= 5 && localY <= TOP_BAR_HEIGHT - 5) {
+                // Click on back button
+                goUpDirectory();
+                return true;
+            }
+        } else if (localY > TOP_BAR_HEIGHT) {
+            const midX = this.size[0] / 2;
+            if (localX < midX) {
+                // Click on directory structure
+                const clickedItem = findClickedItem(directoryStructure, localY - TOP_BAR_HEIGHT + scrollOffsetLeft);
+                if (clickedItem) {
+                    currentDirectory = clickedItem.path;
+                    rootDirectoryWidget.value = currentDirectory;
+                    updateDirectoryStructure();
+                }
+            } else {
+                // Click on thumbnails
+                const thumbnailsPerRow = Math.floor((this.size[0] / 2 - 10) / (THUMBNAIL_SIZE + THUMBNAIL_PADDING));
+                const clickedRow = Math.floor((localY - TOP_BAR_HEIGHT + scrollOffsetRight) / (THUMBNAIL_SIZE + THUMBNAIL_PADDING));
+                const clickedCol = Math.floor((localX - midX) / (THUMBNAIL_SIZE + THUMBNAIL_PADDING));
+                const clickedIndex = clickedRow * thumbnailsPerRow + clickedCol;
+                if (clickedIndex >= 0 && clickedIndex < fileList.length) {
+                    updateSelectedFile(fileList[clickedIndex]);
+                }
+            }
+        }
+        return true;
+    };
+
+    node.onMouseMove = function(event) {
+        if (event.dragging && event.button === 2) { // Right mouse button
+            const midX = this.size[0] / 2;
+            if (event.canvasX - this.pos[0] < midX) {
+                scrollOffsetLeft = Math.max(0, scrollOffsetLeft - event.deltay);
+            } else {
+                scrollOffsetRight = Math.max(0, scrollOffsetRight - event.deltay);
+            }
+            this.setDirtyCanvas(true);
+        }
+    };
+
+    function findClickedItem(structure, y, currentY = 0) {
+        for (const item of structure.children) {
+            if (y >= currentY && y < currentY + FOLDER_HEIGHT) {
+                return item;
+            }
+            currentY += FOLDER_HEIGHT;
+            if (item.children) {
+                const found = findClickedItem(item, y, currentY);
+                if (found) return found;
+                currentY += FOLDER_HEIGHT * countItems(item);
+            }
+        }
+        return null;
+    }
+
+    function countItems(item) {
+        return item.children ? item.children.reduce((sum, child) => sum + countItems(child), 1) : 1;
+    }
+
+    function updateNodeSize() {
+        const width = Math.max(MIN_WIDTH, node.size[0]);
+        const height = Math.max(MIN_HEIGHT, node.size[1]);
+        node.size[0] = width;
+        node.size[1] = height;
+    }
+
+    node.onResize = function() {
+        updateNodeSize();
+        this.setDirtyCanvas(true);
+    };
+
+    updateDirectoryStructure();
+    updateNodeSize();
+}