diff --git a/ComfyUI-Documents/__init__.py b/ComfyUI-Documents/__init__.py new file mode 100644 index 0000000..cde5af3 --- /dev/null +++ b/ComfyUI-Documents/__init__.py @@ -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"]) \ No newline at end of file diff --git a/ComfyUI-Documents/__pycache__/__init__.cpython-311.pyc b/ComfyUI-Documents/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000..6595c7e Binary files /dev/null and b/ComfyUI-Documents/__pycache__/__init__.cpython-311.pyc differ diff --git a/ComfyUI-Documents/__pycache__/comfy_documents_utils.cpython-311.pyc b/ComfyUI-Documents/__pycache__/comfy_documents_utils.cpython-311.pyc new file mode 100644 index 0000000..a7a729d Binary files /dev/null and b/ComfyUI-Documents/__pycache__/comfy_documents_utils.cpython-311.pyc differ diff --git a/ComfyUI-Documents/__pycache__/document_nodes.cpython-311.pyc b/ComfyUI-Documents/__pycache__/document_nodes.cpython-311.pyc new file mode 100644 index 0000000..94556ff Binary files /dev/null and b/ComfyUI-Documents/__pycache__/document_nodes.cpython-311.pyc differ diff --git a/ComfyUI-Documents/__pycache__/server.cpython-311.pyc b/ComfyUI-Documents/__pycache__/server.cpython-311.pyc new file mode 100644 index 0000000..adbda59 Binary files /dev/null and b/ComfyUI-Documents/__pycache__/server.cpython-311.pyc differ diff --git a/ComfyUI-Documents/__pycache__/utils.cpython-311.pyc b/ComfyUI-Documents/__pycache__/utils.cpython-311.pyc new file mode 100644 index 0000000..9f60625 Binary files /dev/null and b/ComfyUI-Documents/__pycache__/utils.cpython-311.pyc differ diff --git a/ComfyUI-Documents/comfy_documents_utils.py b/ComfyUI-Documents/comfy_documents_utils.py new file mode 100644 index 0000000..60bdbc5 --- /dev/null +++ b/ComfyUI-Documents/comfy_documents_utils.py @@ -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) \ No newline at end of file diff --git a/ComfyUI-Documents/document_nodes.py b/ComfyUI-Documents/document_nodes.py new file mode 100644 index 0000000..9857e89 --- /dev/null +++ b/ComfyUI-Documents/document_nodes.py @@ -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" +} \ No newline at end of file diff --git a/ComfyUI-Documents/documents.core.js b/ComfyUI-Documents/documents.core.js new file mode 100644 index 0000000..e7b5406 --- /dev/null +++ b/ComfyUI-Documents/documents.core.js @@ -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 = 'File Path
'; + 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; + }); + } + } +}); \ No newline at end of file diff --git a/ComfyUI-Documents/install.py b/ComfyUI-Documents/install.py new file mode 100644 index 0000000..c7c3778 --- /dev/null +++ b/ComfyUI-Documents/install.py @@ -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() \ No newline at end of file diff --git a/ComfyUI-Documents/requirements.txt b/ComfyUI-Documents/requirements.txt new file mode 100644 index 0000000..61d1cdb --- /dev/null +++ b/ComfyUI-Documents/requirements.txt @@ -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 \ No newline at end of file diff --git a/ComfyUI-Documents/server.py b/ComfyUI-Documents/server.py new file mode 100644 index 0000000..197297e --- /dev/null +++ b/ComfyUI-Documents/server.py @@ -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") \ No newline at end of file diff --git a/ComfyUI-Documents/utils.py b/ComfyUI-Documents/utils.py new file mode 100644 index 0000000..04245d2 --- /dev/null +++ b/ComfyUI-Documents/utils.py @@ -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 \ No newline at end of file diff --git a/ComfyUI-Documents/web/js/documents.core.js b/ComfyUI-Documents/web/js/documents.core.js new file mode 100644 index 0000000..341e66d --- /dev/null +++ b/ComfyUI-Documents/web/js/documents.core.js @@ -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); + }); + } + } +}); \ No newline at end of file