Add files via upload

2024-07-09 19:41:06 +10:00
parent 25d75ba64d
commit 71a44fffc9
14 changed files with 677 additions and 0 deletions
@@ -0,0 +1,44 @@
+import os
+import subprocess
+import sys
+from .document_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+from .utils import folder_paths
+from .server import WEB_DIRECTORY
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']
+
+def install_requirements():
+    requirements_path = os.path.join(os.path.dirname(__file__), "requirements.txt")
+    try:
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", requirements_path])
+        print("Successfully installed requirements for ComfyUI-Documents")
+    except subprocess.CalledProcessError as e:
+        print(f"Error installing requirements: {e}")
+
+try:
+    import fitz  # PyMuPDF
+except ImportError:
+    print("Attempting to install required packages for ComfyUI-Documents...")
+    install_requirements()
+    import fitz
+
+# Ensure the input directory exists
+input_dir = folder_paths.get_input_directory()
+os.makedirs(input_dir, exist_ok=True)
+
+# Set up the web directory for the JavaScript file
+javascript_file = os.path.join(WEB_DIRECTORY, "js", "documents.core.js")
+
+if os.path.exists(javascript_file):
+    print(f"documents.core.js found at {javascript_file}")
+else:
+    print(f"Warning: documents.core.js not found at {javascript_file}")
+
+def get_javascript():
+    if os.path.exists(javascript_file):
+        with open(javascript_file, "r", encoding="utf-8") as file:
+            return file.read()
+    else:
+        return ""
+
+__all__.extend(["get_javascript"])
@@ -0,0 +1,21 @@
+import os
+from typing import List
+
+def get_input_directory():
+    return os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))), "input")
+
+def get_supported_extensions():
+    return [".pdf", ".txt", ".doc", ".docx"]
+
+def get_filename_list() -> List[str]:
+    input_dir = get_input_directory()
+    extensions = get_supported_extensions()
+    
+    files = []
+    for file in os.listdir(input_dir):
+        if os.path.isfile(os.path.join(input_dir, file)):
+            _, ext = os.path.splitext(file)
+            if ext.lower() in extensions:
+                files.append(file)
+    
+    return sorted(files)
@@ -0,0 +1,213 @@
+import os
+import fitz  # PyMuPDF
+import torch
+from PIL import Image
+import numpy as np
+import folder_paths
+from .utils import strip_path
+from docx import Document
+
+class DocumentLoaderNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        input_dir = folder_paths.get_input_directory()
+        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) and f.lower().endswith((".pdf", ".txt", ".doc", ".docx"))]
+        return {"required": {"file_path": (sorted(files),)}}
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("parsed_text",)
+    FUNCTION = "load_document"
+    CATEGORY = "document_processing"
+
+    def load_document(self, file_path):
+        full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
+        _, ext = os.path.splitext(full_path)
+        if ext.lower() == '.pdf':
+            return (self.parse_pdf(full_path),)
+        elif ext.lower() in ['.txt', '.doc', '.docx']:
+            return (self.parse_text(full_path),)
+        else:
+            raise ValueError(f"Unsupported file type: {ext}")
+
+    def parse_pdf(self, file_path):
+        text = ""
+        with fitz.open(file_path) as doc:
+            for page in doc:
+                text += page.get_text()
+        return text
+
+    def parse_text(self, file_path):
+        _, ext = os.path.splitext(file_path)
+        if ext.lower() == '.txt':
+            with open(file_path, 'r', encoding='utf-8') as file:
+                return file.read()
+        elif ext.lower() in ['.doc', '.docx']:
+            doc = Document(file_path)
+            return "\n".join([para.text for para in doc.paragraphs])
+    
+    @classmethod
+    def IS_CHANGED(cls, file_path):
+        full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
+        return full_path
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, file_path):
+        full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
+        return os.path.isfile(full_path)
+
+class PDFToImageNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        input_dir = folder_paths.get_input_directory()
+        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) and f.lower().endswith(".pdf")]
+        return {
+            "required": {
+                "file_path": (sorted(files),),
+                "start_page": ("INT", {"default": 1, "min": 1, "max": 10000}),
+                "end_page": ("INT", {"default": 1, "min": 1, "max": 10000}),
+                "dpi": ("INT", {"default": 300, "min": 72, "max": 600}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "pdf_to_images"
+    CATEGORY = "document_processing"
+    OUTPUT_IS_LIST = (True,)
+
+    def pdf_to_images(self, file_path, start_page, end_page, dpi):
+        full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
+        doc = fitz.open(full_path)
+        num_pages = len(doc)
+        
+        start_page = max(1, min(start_page, num_pages))
+        end_page = max(start_page, min(end_page, num_pages))
+        
+        images = []
+        for page_num in range(start_page - 1, end_page):
+            page = doc.load_page(page_num)
+            pix = page.get_pixmap(matrix=fitz.Matrix(dpi/72, dpi/72))
+            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+            img_np = np.array(img).astype(np.float32) / 255.0
+            img_tensor = torch.from_numpy(img_np)[None,]
+            images.append(img_tensor)
+        
+        doc.close()
+        return (images,)
+    
+    @classmethod
+    def IS_CHANGED(cls, file_path, start_page, end_page, dpi):
+        full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
+        return full_path
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, file_path, start_page, end_page, dpi):
+        full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
+        return os.path.isfile(full_path)
+
+class PDFPageSplitterNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        input_dir = folder_paths.get_input_directory()
+        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) and f.lower().endswith(".pdf")]
+        return {
+            "required": {
+                "file_path": (sorted(files),),
+                "page_numbers": ("STRING", {"default": "1,2,3"}),
+                "dpi": ("INT", {"default": 300, "min": 72, "max": 600}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("selected_pages",)
+    FUNCTION = "pdf_to_selected_images"
+    CATEGORY = "document_processing"
+    OUTPUT_NODE = True
+    OUTPUT_IS_LIST = (True,)
+
+    def pdf_to_selected_images(self, file_path, page_numbers, dpi):
+        full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
+        doc = fitz.open(full_path)
+        num_pages = len(doc)
+        
+        # Parse page numbers from input string
+        try:
+            selected_pages = [int(p.strip()) for p in page_numbers.split(',') if p.strip()]
+            selected_pages = [p for p in selected_pages if 1 <= p <= num_pages]
+        except ValueError:
+            raise ValueError("Invalid page numbers. Please provide comma-separated integers.")
+        
+        images = []
+        for page_num in selected_pages:
+            page = doc.load_page(page_num - 1)  # PDF pages are 0-indexed
+            pix = page.get_pixmap(matrix=fitz.Matrix(dpi/72, dpi/72))
+            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+            img_np = np.array(img).astype(np.float32) / 255.0
+            img_tensor = torch.from_numpy(img_np)[None,]
+            images.append(img_tensor)
+        
+        doc.close()
+        return (images,)
+
+    @classmethod
+    def IS_CHANGED(cls, file_path, page_numbers, dpi):
+        full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
+        return full_path
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, file_path, page_numbers, dpi):
+        full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
+        if not os.path.isfile(full_path):
+            return False
+        try:
+            [int(p.strip()) for p in page_numbers.split(',') if p.strip()]
+        except ValueError:
+            return False
+        return True
+
+class ImageSelectorNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "indexes": ("STRING", {"default": "0"}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "select_images"
+    CATEGORY = "image_processing"
+    OUTPUT_NODE = True
+    OUTPUT_IS_LIST = (True,)
+
+    def select_images(self, images, indexes):
+        try:
+            index_list = [int(idx.strip()) for idx in indexes.split(',')]
+        except ValueError:
+            raise ValueError("Invalid indexes. Please provide comma-separated integers.")
+        
+        selected_images = []
+        for idx in index_list:
+            if idx < 0 or idx >= len(images):
+                raise ValueError(f"Index {idx} is out of range. Available images: {len(images)}")
+            selected_images.append(images[idx])
+        
+        return (selected_images,)
+
+    @classmethod
+    def IS_CHANGED(cls, images, indexes):
+        return indexes
+
+NODE_CLASS_MAPPINGS = {
+    "DocumentLoader": DocumentLoaderNode,
+    "PDFToImage": PDFToImageNode,
+    "PDFPageSplitter": PDFPageSplitterNode,
+    "ImageSelector": ImageSelectorNode
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "DocumentLoader": "Document Loader",
+    "PDFToImage": "PDF to Image (Multi-Page)",
+    "PDFPageSplitter": "PDF Page Splitter",
+    "ImageSelector": "Image Selector"
+}
@@ -0,0 +1,160 @@
+import { app } from '../../../scripts/app.js';
+import { api } from '../../../scripts/api.js';
+
+function chainCallback(object, property, callback) {
+    if (object == undefined) {
+        console.error("Tried to add callback to non-existent object");
+        return;
+    }
+    if (property in object) {
+        const callback_orig = object[property];
+        object[property] = function () {
+            const r = callback_orig.apply(this, arguments);
+            callback.apply(this, arguments);
+            return r;
+        };
+    } else {
+        object[property] = callback;
+    }
+}
+
+function searchBox(event, [x, y], node) {
+    if (this.prompt) {
+        return;
+    }
+    this.prompt = true;
+
+    let pathWidget = this;
+    let dialog = document.createElement("div");
+    dialog.className = "litegraph litesearchbox graphdialog rounded";
+    dialog.innerHTML = '<span class="name">File Path</span> <input autofocus="" type="text" class="value"><button class="rounded">OK</button><div class="helper"></div>';
+    dialog.close = () => {
+        dialog.remove();
+    };
+    document.body.append(dialog);
+    if (app.canvas.ds.scale > 1) {
+        dialog.style.transform = "scale(" + app.canvas.ds.scale + ")";
+    }
+    var input = dialog.querySelector(".value");
+    var options_element = dialog.querySelector(".helper");
+    input.value = pathWidget.value || "";
+
+    var timeout = null;
+    let last_path = null;
+    let extensions = pathWidget.options?.extensions;
+
+    input.addEventListener("keydown", (e) => {
+        dialog.is_modified = true;
+        if (e.keyCode == 27) {
+            dialog.close();
+        } else if (e.keyCode == 13 && e.target.localName != "textarea") {
+            if (pathWidget.value) {
+                pathWidget.value = input.value;
+                if (pathWidget.callback) {
+                    pathWidget.callback(pathWidget.value);
+                }
+            }
+            dialog.close();
+        } else {
+            if (timeout) {
+                clearTimeout(timeout);
+            }
+            timeout = setTimeout(updateOptions, 100);
+            return;
+        }
+        this.prompt = false;
+        e.preventDefault();
+        e.stopPropagation();
+    });
+
+    var button = dialog.querySelector("button");
+    button.addEventListener("click", (e) => {
+        if (pathWidget.value) {
+            pathWidget.value = input.value;
+            if (pathWidget.callback) {
+                pathWidget.callback(pathWidget.value);
+            }
+        }
+        node.graph.setDirtyCanvas(true);
+        dialog.close();
+        this.prompt = false;
+    });
+
+    var rect = app.canvas.canvas.getBoundingClientRect();
+    var offsetx = -20;
+    var offsety = -20;
+    if (rect) {
+        offsetx -= rect.left;
+        offsety -= rect.top;
+    }
+
+    if (event) {
+        dialog.style.left = event.clientX + offsetx + "px";
+        dialog.style.top = event.clientY + offsety + "px";
+    } else {
+        dialog.style.left = canvas.width * 0.5 + offsetx + "px";
+        dialog.style.top = canvas.height * 0.5 + offsety + "px";
+    }
+
+    async function updateOptions() {
+        timeout = null;
+        let params = { path: input.value };
+        if (extensions) {
+            params.extensions = extensions;
+        }
+        let optionsURL = api.apiURL('getpath?' + new URLSearchParams(params));
+        try {
+            let resp = await fetch(optionsURL);
+            let options = await resp.json();
+            options_element.innerHTML = '';
+            for (let option of options) {
+                let el = document.createElement("div");
+                el.innerText = option;
+                el.className = "litegraph lite-search-item";
+                if (option.endsWith('/')) {
+                    el.className += " is-dir";
+                    el.addEventListener("click", (e) => {
+                        input.value = option;
+                        if (timeout) {
+                            clearTimeout(timeout);
+                        }
+                        timeout = setTimeout(updateOptions, 100);
+                    });
+                } else {
+                    el.addEventListener("click", (e) => {
+                        if (pathWidget.value) {
+                            pathWidget.value = option;
+                            if (pathWidget.callback) {
+                                pathWidget.callback(pathWidget.value);
+                            }
+                        }
+                        dialog.close();
+                        pathWidget.prompt = false;
+                    });
+                }
+                options_element.appendChild(el);
+            }
+        } catch (e) {
+            console.error("Error fetching options:", e);
+        }
+    }
+
+    setTimeout(async function () {
+        input.focus();
+        await updateOptions();
+    }, 100);
+
+    return dialog;
+}
+
+app.registerExtension({
+    name: "ComfyUI-Documents",
+    async beforeRegisterNodeDef(nodeType, nodeData, app) {
+        if (nodeData?.name == "DocumentLoader" || nodeData?.name == "PDFToImage") {
+            chainCallback(nodeType.prototype, "onNodeCreated", function () {
+                const pathWidget = this.widgets.find((w) => w.name === "document");
+                pathWidget.mouse = searchBox;
+            });
+        }
+    }
+});
@@ -0,0 +1,14 @@
+import subprocess
+import sys
+import os
+
+def install_requirements():
+    requirements_path = os.path.join(os.path.dirname(__file__), "requirements.txt")
+    try:
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", requirements_path])
+        print("Successfully installed requirements for ComfyUI-Documents")
+    except subprocess.CalledProcessError as e:
+        print(f"Error installing requirements: {e}")
+        
+if __name__ == "__main__":
+    install_requirements()
@@ -0,0 +1,5 @@
+PyMuPDF==1.23.8
+python-docx==1.1.0
+pandas==2.1.4
+Markdown==3.5.1
+PyYAML==6.0.1
@@ -0,0 +1,55 @@
+import server
+import os
+from aiohttp import web
+from .utils import is_safe_path, strip_path
+from .utils import folder_paths  # Import folder_paths from utils
+
+routes = server.PromptServer.instance.routes
+
+@routes.post("/upload/document")
+async def upload_document(request):
+    try:
+        data = await request.post()
+        file = data['document']
+        filename = file.filename
+        content = file.file.read()
+        
+        input_dir = folder_paths.get_input_directory()
+        file_path = os.path.join(input_dir, filename)
+        
+        if not is_safe_path(file_path):
+            return web.Response(status=403, text="Invalid file path")
+        
+        with open(file_path, 'wb') as f:
+            f.write(content)
+        
+        return web.json_response({"name": filename})
+    except Exception as e:
+        print(f"Error in upload_document: {str(e)}")
+        return web.Response(status=500, text=str(e))
+
+@routes.get("/getpath")
+async def get_path(request):
+    query = request.rel_url.query
+    if "path" not in query:
+        return web.Response(status=400, text="Missing 'path' parameter")
+    path = os.path.abspath(strip_path(query["path"]))
+
+    if not os.path.exists(path) or not is_safe_path(path):
+        return web.json_response([])
+
+    valid_extensions = query.get("extensions", "").split(",")
+    valid_items = []
+    for item in os.scandir(path):
+        try:
+            if item.is_dir():
+                valid_items.append(item.name + "/")
+                continue
+            if not valid_extensions or item.name.split(".")[-1].lower() in valid_extensions:
+                valid_items.append(item.name)
+        except OSError:
+            pass
+
+    return web.json_response(valid_items)
+
+WEB_DIRECTORY = os.path.join(os.path.dirname(__file__), "web")
@@ -0,0 +1,36 @@
+import os
+
+class folder_paths:
+    @classmethod
+    def get_input_directory(cls):
+        # Update the path to the correct location of the 'input' directory
+        return os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))), "input")
+
+    @classmethod
+    def get_annotated_filepath(cls, filename):
+        input_dir = cls.get_input_directory()
+        file_path = os.path.join(input_dir, filename)
+        return file_path, input_dir
+
+    @classmethod
+    def exists_annotated_filepath(cls, filename):
+        file_path, _ = cls.get_annotated_filepath(filename)
+        return os.path.exists(file_path)
+
+def strip_path(path):
+    path = path.strip()
+    if path.startswith("\""):
+        path = path[1:]
+    if path.endswith("\""):
+        path = path[:-1]
+    return path
+
+def is_safe_path(path):
+    if "STRICT_PATHS" not in os.environ:
+        return True
+    basedir = os.path.abspath('.')
+    try:
+        common_path = os.path.commonpath([basedir, path])
+    except:
+        return False
+    return common_path == basedir
@@ -0,0 +1,129 @@
+import { app } from '../../../scripts/app.js';
+import { api } from '../../../scripts/api.js';
+
+function chainCallback(object, property, callback) {
+    if (object == undefined) {
+        console.error("Tried to add callback to non-existent object");
+        return;
+    }
+    if (property in object) {
+        const callback_orig = object[property];
+        object[property] = function () {
+            const r = callback_orig.apply(this, arguments);
+            callback.apply(this, arguments);
+            return r;
+        };
+    } else {
+        object[property] = callback;
+    }
+}
+
+function createFileUploadWidget(node, inputName, inputData, app) {
+    const fileInput = document.createElement("input");
+    Object.assign(fileInput, {
+        type: "file",
+        accept: ".pdf,.txt,.doc,.docx",
+        style: "display: none",
+        onchange: async () => {
+            if (fileInput.files.length) {
+                await uploadFile(fileInput.files[0], true);
+            }
+        },
+    });
+    document.body.append(fileInput);
+
+    const uploadWidget = node.addWidget("button", inputName, "document", () => {
+        fileInput.click();
+    });
+    uploadWidget.label = "Choose file to upload";
+    uploadWidget.serialize = false;
+
+    async function uploadFile(file, updateNode) {
+        try {
+            const body = new FormData();
+            body.append("document", file);
+            const resp = await api.fetchApi("/upload/document", {
+                method: "POST",
+                body,
+            });
+            if (resp.status === 200) {
+                const data = await resp.json();
+                let path = data.name;
+                if (data.subfolder) path = data.subfolder + "/" + path;
+                
+                let documentWidget = null;
+                for (let i = 0; i < node.widgets.length; i++) {
+                    if (node.widgets[i].name === "file_path") {
+                        documentWidget = node.widgets[i];
+                        break;
+                    }
+                }
+                
+                if (documentWidget && !documentWidget.options.values.includes(path)) {
+                    documentWidget.options.values.push(path);
+                }
+                if (updateNode && documentWidget) {
+                    documentWidget.value = path;
+                }
+            } else {
+                alert(resp.status + " - " + resp.statusText);
+            }
+        } catch (error) {
+            alert(error);
+        }
+    }
+
+    node.onDragOver = function (e) {
+        if (e.dataTransfer && e.dataTransfer.items) {
+            const document = [...e.dataTransfer.items].find((f) => f.kind === "file" && 
+                (f.type === "application/pdf" || f.type === "text/plain" || 
+                 f.type === "application/msword" || f.type === "application/vnd.openxmlformats-officedocument.wordprocessingml.document"));
+            return !!document;
+        }
+        return false;
+    };
+
+    node.onDragDrop = function (e) {
+        let handled = false;
+        for (const file of e.dataTransfer.files) {
+            if (file.type === "application/pdf" || file.type === "text/plain" || 
+                file.type === "application/msword" || file.type === "application/vnd.openxmlformats-officedocument.wordprocessingml.document") {
+                uploadFile(file, !handled);
+                handled = true;
+            }
+        }
+        return handled;
+    };
+
+    return { widget: uploadWidget };
+}
+
+app.registerExtension({
+    name: "ComfyUI-Documents",
+    async beforeRegisterNodeDef(nodeType, nodeData, app) {
+        if (nodeData?.name == "DocumentLoader" || nodeData?.name == "PDFToImage") {
+            chainCallback(nodeType.prototype, "onNodeCreated", function () {
+                let documentWidget = null;
+                for (let i = 0; i < this.widgets.length; i++) {
+                    if (this.widgets[i].name === "file_path") {
+                        documentWidget = this.widgets[i];
+                        break;
+                    }
+                }
+                
+                if (documentWidget) {
+                    documentWidget.options.values = [];
+                    api.fetchApi('/getpath?path=input&extensions=pdf,txt,doc,docx')
+                        .then(response => response.json())
+                        .then(files => {
+                            documentWidget.options.values = files;
+                            if (files.length > 0) {
+                                documentWidget.value = files[0];
+                            }
+                        });
+                }
+                createFileUploadWidget(this, "upload_document", {}, app);
+            });
+        }
+    }
+});