From 4189bc62ea805fcb7fdd0da4596c025ed2917e36 Mon Sep 17 00:00:00 2001 From: S4MUEL-404 Date: Thu, 18 Sep 2025 17:27:52 +0800 Subject: [PATCH] Add new features and update documentation - Add WARP.md documentation - Add saveByFileName.py node for file saving functionality - Update README files with latest version info - Update __init__.py and requirements.txt --- README.md | 7 +- README_ZHS.md | 7 +- README_ZHT.md | 7 +- WARP.md | 113 ++++++++ __init__.py | 3 + py/saveByFileName.py | 599 +++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 8 +- 7 files changed, 739 insertions(+), 5 deletions(-) create mode 100644 WARP.md create mode 100644 py/saveByFileName.py diff --git a/README.md b/README.md index a86fbaa..b7a617f 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,9 @@ A comprehensive workflow optimization toolkit for ComfyUI, providing essential n - **💀Prepack Int Combine** - Combine up to 4 integers into a string with selectable separator - **💀Prepack Int Split** - Split a string into up to 4 integers using selectable separator +### File Management +- **💀Save By File Name** - Smart file saving with format preservation and custom naming. Supports images (WebP, JPEG, PNG, GIF), videos (MP4, AVI), and text files with automatic format detection + ## 📦 Installation ### Method 1: ComfyUI Manager (Recommended) @@ -81,6 +84,7 @@ All dependencies are typically already available in standard ComfyUI installatio - ✅ **Pipeline Management** - Store and retrieve workflow states efficiently - ✅ **Seed Management** - Smart seed control with history tracking - ✅ **Logic Operations** - Comprehensive logic and comparison tools +- ✅ **File Management** - Smart file saving with format preservation and custom naming ## 📁 Project Structure @@ -98,7 +102,8 @@ ComfyUI-Prepack/ │ ├── logicInt.py # Integer logic │ ├── logicString.py # String logic │ ├── intCombine.py # Integer combination -│ └── intSplit.py # Integer splitting +│ ├── intSplit.py # Integer splitting +│ └── saveByFileName.py # Smart file saving ├── js/ # JavaScript UI extensions │ ├── seed.js # Seed management UI │ ├── loraText.js # LoRA text integration diff --git a/README_ZHS.md b/README_ZHS.md index 0b7891d..42603db 100644 --- a/README_ZHS.md +++ b/README_ZHS.md @@ -29,6 +29,9 @@ - **💀Prepack Int Combine** - 将最多 4 个整数合并为字符串，可选分隔符 - **💀Prepack Int Split** - 使用可选分隔符将字符串拆分为最多 4 个整数 +### 文件管理 +- **💀Save By File Name** - 智能文件保存，具备格式保持和自定义命名功能。支持图片（WebP、JPEG、PNG、GIF）、视频（MP4、AVI）和文本文件，具备自动格式检测功能 + ## 📦 安装方法 ### 方法一：ComfyUI Manager（推荐） @@ -81,6 +84,7 @@ - ✅ **管道管理** - 高效存储和获取工作流状态 - ✅ **种子管理** - 智能种子控制，具备历史跟踪 - ✅ **逻辑运算** - 综合逻辑和比较工具 +- ✅ **文件管理** - 智能文件保存，具备格式保持和自定义命名功能 ## 📁 项目结构 @@ -98,7 +102,8 @@ ComfyUI-Prepack/ │ ├── logicInt.py # 整数逻辑 │ ├── logicString.py # 字符串逻辑 │ ├── intCombine.py # 整数合并 -│ └── intSplit.py # 整数拆分 +│ ├── intSplit.py # 整数拆分 +│ └── saveByFileName.py # 智能文件保存 ├── js/ # JavaScript UI 扩展 │ ├── seed.js # 种子管理 UI │ ├── loraText.js # LoRA 文本集成 diff --git a/README_ZHT.md b/README_ZHT.md index 249dacc..997f50a 100644 --- a/README_ZHT.md +++ b/README_ZHT.md @@ -29,6 +29,9 @@ - **💀Prepack Int Combine** - 將最多 4 個整數合併為字串，可選分隔符 - **💀Prepack Int Split** - 使用可選分隔符將字串拆分為最多 4 個整數 +### 文件管理 +- **💀Save By File Name** - 智慧文件保存，具備格式保持和自訂命名功能。支援圖片（WebP、JPEG、PNG、GIF）、影片（MP4、AVI）和文本文件，具備自動格式檢測功能 + ## 📦 安裝方法 ### 方法一：ComfyUI Manager（推薦） @@ -81,6 +84,7 @@ - ✅ **管道管理** - 高效率儲存和取得工作流狀態 - ✅ **種子管理** - 智慧種子控制，具備歷史追蹤 - ✅ **邏輯運算** - 綜合邏輯和比較工具 +- ✅ **文件管理** - 智慧文件保存，具備格式保持和自訂命名功能 ## 📁 專案結構 @@ -98,7 +102,8 @@ ComfyUI-Prepack/ │ ├── logicInt.py # 整數邏輯 │ ├── logicString.py # 字串邏輯 │ ├── intCombine.py # 整數合併 -│ └── intSplit.py # 整數拆分 +│ ├── intSplit.py # 整數拆分 +│ └── saveByFileName.py # 智慧文件保存 ├── js/ # JavaScript UI 擴展 │ ├── seed.js # 種子管理 UI │ ├── loraText.js # LoRA 文本整合 diff --git a/WARP.md b/WARP.md new file mode 100644 index 0000000..ecb8b54 --- /dev/null +++ b/WARP.md @@ -0,0 +1,113 @@ +# WARP.md + +這個文件為 WARP (warp.dev) 在這個倉庫中工作時提供指導。 + +## 專案概述 + +ComfyUI-Prepack 是一個為 ComfyUI 設計的綜合工作流優化工具包。它提供了一套帶有 💀 前綴的自定義節點，涵蓋模型管理、採樣控制、工作流管理和邏輯運算等核心功能。這個專案的特點是將複雜的 ComfyUI 操作簡化為單個節點，提升工作流效率。 + +## 核心架構 + +### 模組結構 +- **py/** - Python 後端節點實現，每個檔案對應一個具體的節點功能 +- **js/** - JavaScript 前端 UI 擴展，增強使用者界面體驗 +- **__init__.py** - 主入口點，定義所有節點映射和分類 + +### 節點類別 +1. **模型管理** - PrepackModelDualCLIP/SingleCLIP (支援多種 CLIP 類型：SDXL、SD3、FLUX、Hunyuan Video) +2. **LoRA 管理** - PrepackLoras/LorasAndMSSD3 (支援最多3個 LoRA 和文本文件整合) +3. **採樣控制** - PrepackKsampler/KsamplerAdvanced (具備完整錯誤處理和調試功能) +4. **工作流管理** - PrepackSetPipe/GetPipe (管道狀態的存儲和檢索) +5. **智能種子** - PrepackSeed (帶有歷史記錄和隨機生成按鈕) +6. **邏輯運算** - PrepackLogicInt/String 和 PrepackIntCombine/Split +7. **文件保存** - PrepackSaveByFileName (支持圖片、視頻、文本的自定義文件名保存) + +### JavaScript UI 架構 +- **seed.js** - 提供種子歷史追蹤和隨機生成按鈕（最多50條歷史記錄） +- **loraText.js** - LoRA 文本文件整合，支援動態載入 .txt 檔案內容 +- **setgetnodes.js** - Set/Get 虛擬節點系統，具備類型適配和顏色管理 + +## 常用開發命令 + +### 測試與調試 +```bash +# 啟用 Prepack 調試模式 +set PREPACK_DEBUG=1 + +# 查看 ComfyUI 日誌 +# ComfyUI 控制台會顯示節點執行和錯誤資訊 + +# 檢查 LoRA 文本文件 API +# 瀏覽器訪問: http://localhost:8188/prepack/lora-texts/{lora_name} +``` + +### 安裝與部署 +```bash +# 透過 ComfyUI Manager 安裝（推薦） +# 在 ComfyUI Manager 中搜尋 "Prepack" 並安裝 + +# 手動安裝 +cd ComfyUI/custom_nodes/ +git clone https://github.com/S4MUEL-404/ComfyUI-Prepack.git +pip install -r ComfyUI-Prepack/requirements.txt +``` + +## 重要開發模式 + +### 節點開發模式 +- 所有節點都繼承標準 ComfyUI 節點結構 +- 使用 `INPUT_TYPES` 定義輸入，`RETURN_TYPES` 定義輸出 +- 必須提供完整的 `tooltip` 和錯誤處理 +- 節點分類統一使用 `"💀Prepack"` + +### JavaScript 擴展模式 +- 使用 `app.registerExtension` 註冊前端擴展 +- 透過 `beforeRegisterNodeDef` 修改節點行為 +- 使用 ComfyUI 標準 API (`api.fetchApi`) 進行後端通信 +- 所有 UI 修改都應保持 ComfyUI 的原生風格 + +### 管道（Pipe）系統 +PrepackSetPipe 和 PrepackGetPipe 實現了狀態傳遞機制： +- SetPipe 將多個組件打包為單一管道對象 +- GetPipe 解包管道對象為個別組件 +- 支援：model, clip, vae, lora_path, lora_text, positive, negative, latent_image, seed, steps, cfg, denoise + +### LoRA 整合系統 +- 支援最多3個 LoRA 同時載入 +- 自動搜尋同名資料夾中的 .txt 檔案 +- 提供 HTTP API 端點用於動態載入文本內容 +- 安全檢查防止路徑遍歷攻擊 + +### 文件保存系統 +- 多格式支持：圖片 (PNG/JPG/WebP)、視頻 (MP4/AVI/MOV)、文本 (TXT/JSON/CSV) +- 智能文件名處理：支持 {date}、{time}、{timestamp} 佔位符 +- 批量處理和防重複機制 +- 元數據嵌入和質量控制 + +## 關鍵技術實現 + +### 類型適配系統 +Set/Get 節點實現動態類型適配： +- 根據連接自動推斷和適配類型 +- 支援萬用字元 (*) 類型 +- 顏色編碼區分不同資料類型 + +### 種子管理系統 +- 支援64位無符號整數範圍 +- JavaScript 擴展提供歷史追蹤（最多50條） +- 自動監控工作流執行中的種子變化 + +### 錯誤處理模式 +所有節點都實現了完整的錯誤處理： +- 輸入驗證和類型檢查 +- 資源載入失敗的優雅降級 +- 詳細的錯誤訊息和日誌記錄 + +## 專案規則遵循 + +- 所有註釋必須使用英文 +- 程式碼命名遵循既有規則 +- requirements.txt 保持最簡依賴 +- summary_md 目錄用於存放總結文件 +- 主頁地址固定為 https://github.com/S4MUEL-404/ +- 版本資訊和作者署名：S4MUEL (s4muel.com) \ No newline at end of file diff --git a/__init__.py b/__init__.py index b0beb41..efd7856 100644 --- a/__init__.py +++ b/__init__.py @@ -11,6 +11,7 @@ from .py.logicInt import PrepackLogicInt from .py.logicString import PrepackLogicString from .py.intCombine import PrepackIntCombine from .py.intSplit import PrepackIntSplit +from .py.saveByFileName import PrepackSaveByFileName # Frontend extension directory for virtual nodes WEB_DIRECTORY = "./js" @@ -29,6 +30,7 @@ NODE_CLASS_MAPPINGS = { "💀Prepack Logic String": PrepackLogicString, "💀Prepack Int Combine": PrepackIntCombine, "💀Prepack Int Split": PrepackIntSplit, + "💀Prepack Save By File Name": PrepackSaveByFileName, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -45,6 +47,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "💀Prepack Logic String": "💀Prepack Logic String", "💀Prepack Int Combine": "💀Prepack Int Combine", "💀Prepack Int Split": "💀Prepack Int Split", + "💀Prepack Save By File Name": "💀Save By File Name", } __all__ = [ diff --git a/py/saveByFileName.py b/py/saveByFileName.py new file mode 100644 index 0000000..b0ea3f5 --- /dev/null +++ b/py/saveByFileName.py @@ -0,0 +1,599 @@ +import os +import datetime +import shutil +import folder_paths + +"""Prepack Save By File Name: rename and copy files with custom file names without any modification.""" + + +class PrepackSaveByFileName: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + self.type = "output" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "filename": ("STRING", { + "default": "output", + "multiline": False, + "tooltip": "Base filename. Add .webp, .jpg, .png etc. to force specific format. Supports {date}, {time}, {timestamp} placeholders." + }), + "overwrite": (["false", "true"], { + "default": "false", + "tooltip": "Whether to overwrite existing files or add suffix." + }), + }, + "optional": { + "image": ("IMAGE", { + "tooltip": "Image files: PNG, JPG, JPEG, GIF, WebP, APNG formats (preserves animation)." + }), + "video": ("*", { + "tooltip": "Video files: MP4, AVI formats." + }), + "text": ("STRING", { + "tooltip": "Text content to save as file.", + "forceInput": True + }) + } + } + + RETURN_TYPES = ("STRING", "STRING") + RETURN_NAMES = ("file_path", "filename") + OUTPUT_TOOLTIPS = ( + "Full path to the renamed file.", + "Final filename used for renaming." + ) + FUNCTION = "save_by_filename" + OUTPUT_NODE = True + + CATEGORY = "💀Prepack" + DESCRIPTION = "Rename and copy files with custom file names without any modification. Preserves original format and content." + + def save_by_filename(self, filename, overwrite, image=None, video=None, text=None): + # SaveByFileName v1.2 - Format preservation enabled by default + try: + # Find which input was provided + file_data = None + file_type = None + + if image is not None: + file_data = image + file_type = 'image' + elif video is not None: + file_data = video + file_type = 'video' + elif text is not None: + file_data = text + file_type = 'text' + else: + raise ValueError("No input provided. Please connect image, video, or text.") + + # Process filename placeholders + processed_filename = self.process_filename_placeholders(filename) + + # Handle different file types + if file_type == 'text': + # For text input, save directly as text file + output_path = self.determine_output_path(processed_filename, None, 'txt', overwrite) + + with open(output_path, 'w', encoding='utf-8') as f: + f.write(str(file_data)) + + output_filename = os.path.basename(output_path) + print(f"Text saved: {output_path}") + + else: + # For image/video types, try to get source file path first + source_path = self.get_source_file_path(file_data) + + if source_path and os.path.isfile(source_path): + # Found existing file - copy it directly + original_ext = os.path.splitext(source_path)[1].lstrip('.') + if not original_ext: + original_ext = 'png' if file_type == 'image' else 'mp4' + + output_path = self.determine_output_path(processed_filename, source_path, original_ext, overwrite) + shutil.copy2(source_path, output_path) + + output_filename = os.path.basename(output_path) + print(f"File renamed and copied: {source_path} -> {output_path}") + + else: + # No source file found - handle tensor data + if file_type == 'image': + + # For image tensor, try to preserve format if specified by user or detect from tensor + user_ext = None + if '.' in processed_filename: + user_ext = os.path.splitext(processed_filename)[1].lstrip('.') + + # Try to detect original format from tensor metadata + detected_ext = self.detect_image_format(file_data) + + # If no format detected, try to infer from context + if not detected_ext and not user_ext: + detected_ext = self.infer_format_from_context() + + # Priority: user specified > detected format > png default + default_ext = user_ext if user_ext else (detected_ext if detected_ext else 'png') + output_path = self.determine_output_path(processed_filename, None, default_ext, overwrite) + self.save_image_tensor(file_data, output_path) + + elif file_type == 'video': + + # Handle user-specified extension + user_ext = None + if '.' in processed_filename: + user_ext = os.path.splitext(processed_filename)[1].lstrip('.') + + # For video, check if it's a file path or tensor + if isinstance(file_data, str) and os.path.isfile(file_data): + # It's a file path + original_ext = os.path.splitext(file_data)[1].lstrip('.') + if not original_ext: + original_ext = 'mp4' + final_ext = user_ext if user_ext else original_ext + output_path = self.determine_output_path(processed_filename, file_data, final_ext, overwrite) + shutil.copy2(file_data, output_path) + else: + # Handle video objects or tensor data + default_ext = user_ext if user_ext else 'mp4' + output_path = self.determine_output_path(processed_filename, None, default_ext, overwrite) + + # Try to save video data + try: + self.save_video_data(file_data, output_path) + except Exception as save_error: + print(f"Error saving video data: {str(save_error)}") + # Fallback: try to extract file path from video object + video_source = self.extract_video_source_path(file_data) + if video_source and os.path.isfile(video_source): + shutil.copy2(video_source, output_path) + else: + raise Exception(f"Cannot process video data of type {type(file_data)}") + + output_filename = os.path.basename(output_path) + print(f"{file_type.capitalize()} saved: {output_path}") + + return (output_path, output_filename) + + except Exception as e: + print(f"Error in PrepackSaveByFileName: {str(e)}") + return ("", "") + + def process_filename_placeholders(self, filename): + """Process placeholders in filename like {date}, {time}, {timestamp}""" + now = datetime.datetime.now() + + placeholders = { + "{date}": now.strftime("%Y-%m-%d"), + "{time}": now.strftime("%H-%M-%S"), + "{timestamp}": str(int(now.timestamp())), + "{datetime}": now.strftime("%Y-%m-%d_%H-%M-%S") + } + + processed = str(filename) + for placeholder, value in placeholders.items(): + processed = processed.replace(placeholder, value) + + # Remove invalid filename characters + invalid_chars = '<>:"/\\|?*' + for char in invalid_chars: + processed = processed.replace(char, "_") + + return processed + + + def determine_output_path(self, processed_filename, source_path, default_ext, overwrite): + """Determine output path with smart extension handling""" + # Check if filename already has an extension + filename_name, filename_ext = os.path.splitext(processed_filename) + + if filename_ext: # User specified extension in filename + # Use user-specified extension, remove the dot + final_ext = filename_ext.lstrip('.') + final_filename = processed_filename + else: # No extension in filename + # Use source file extension or default + final_ext = default_ext + final_filename = f"{processed_filename}.{final_ext}" + + # Create full output path + output_path = os.path.join(self.output_dir, final_filename) + + # Handle file conflicts + return self.get_unique_filename(output_path, overwrite) + + def get_unique_filename(self, base_path, overwrite): + """Get unique filename if file exists and overwrite is false""" + if overwrite == "true" or not os.path.exists(base_path): + return base_path + + directory = os.path.dirname(base_path) + filename = os.path.basename(base_path) + name, ext = os.path.splitext(filename) + + counter = 1 + while True: + new_filename = f"{name}_{counter:03d}{ext}" + new_path = os.path.join(directory, new_filename) + if not os.path.exists(new_path): + return new_path + counter += 1 + + def get_source_file_path(self, file): + """Get source file path from various input types""" + try: + # Direct string file path + if isinstance(file, str): + if os.path.isfile(file): + return file + # Try to decode if it looks like a path + if '\\' in file or '/' in file: + cleaned_path = file.strip('"\'') + if os.path.isfile(cleaned_path): + return cleaned_path + + # Dictionary with file information + if isinstance(file, dict): + # Common ComfyUI file dict keys + for key in ['filename', 'path', 'file_path', 'filepath', 'source_path', 'src_path', 'source_file', 'original_file']: + if key in file and isinstance(file[key], str) and os.path.isfile(file[key]): + return file[key] + + # Check for nested dictionaries + for key, value in file.items(): + if isinstance(value, (dict, str)): + nested_path = self.get_source_file_path(value) + if nested_path: + return nested_path + + # List or tuple - check all elements + if isinstance(file, (list, tuple)): + for item in file: + path = self.get_source_file_path(item) + if path: + return path + + # Try hasattr for objects with file attributes + if hasattr(file, '__dict__'): + for attr_name in ['filename', 'path', 'file_path', 'source', 'source_file', 'original_file']: + if hasattr(file, attr_name): + attr_value = getattr(file, attr_name) + if isinstance(attr_value, str) and os.path.isfile(attr_value): + return attr_value + + return None + + except Exception as e: + print(f"Error getting source file path: {str(e)}") + return None + + def detect_image_format(self, file_data): + """Try to detect image format from various sources""" + try: + # Check if file_data has format information + if hasattr(file_data, 'format') and file_data.format: + format_name = str(file_data.format).lower() + if format_name in ['jpeg', 'jpg', 'png', 'gif', 'webp', 'bmp']: + return 'jpg' if format_name == 'jpeg' else format_name + + # Check metadata or attributes + if hasattr(file_data, '__dict__'): + for attr in ['format', 'format_name', 'file_format', 'extension', 'ext', 'source', 'filename']: + if hasattr(file_data, attr): + value = str(getattr(file_data, attr)).lower() + if value in ['jpeg', 'jpg', 'png', 'gif', 'webp', 'bmp']: + return 'jpg' if value == 'jpeg' else value + # Check if it's a path + if '.' in value: + ext = os.path.splitext(value)[1].lower().lstrip('.') + if ext in ['jpg', 'jpeg', 'png', 'gif', 'webp', 'bmp']: + return 'jpg' if ext == 'jpeg' else ext + + # Check if it's a dictionary with format info + if isinstance(file_data, dict): + for key in ['format', 'file_format', 'extension', 'ext', 'type', 'source', 'filename']: + if key in file_data: + value = str(file_data[key]).lower() + if value in ['jpeg', 'jpg', 'png', 'gif', 'webp', 'bmp']: + return 'jpg' if value == 'jpeg' else value + # Check if it's a path + if '.' in value: + ext = os.path.splitext(value)[1].lower().lstrip('.') + if ext in ['jpg', 'jpeg', 'png', 'gif', 'webp', 'bmp']: + return 'jpg' if ext == 'jpeg' else ext + + # Try to get format from potential file path in data + potential_path = self.get_source_file_path(file_data) + if potential_path: + ext = os.path.splitext(potential_path)[1].lower().lstrip('.') + if ext in ['jpg', 'jpeg', 'png', 'gif', 'webp', 'bmp']: + return 'jpg' if ext == 'jpeg' else ext + + return None + + except Exception as e: + return None + + def infer_format_from_context(self): + """Try to infer image format from context clues""" + try: + import folder_paths + input_dir = folder_paths.get_input_directory() + if os.path.isdir(input_dir): + current_time = datetime.datetime.now().timestamp() + for file in os.listdir(input_dir): + if file.lower().endswith('.webp'): + file_path = os.path.join(input_dir, file) + if os.path.isfile(file_path): + mod_time = os.path.getmtime(file_path) + if current_time - mod_time < 300: # Within 5 minutes + return 'webp' + return None + except Exception as e: + return None + + def extract_video_source_path(self, video_data): + """Extract source file path from video object""" + try: + # Check for methods that might return source information + method_attrs = ['get_stream_source', 'get_source_path', 'get_file_path'] + for method_name in method_attrs: + if hasattr(video_data, method_name): + try: + method = getattr(video_data, method_name) + if callable(method): + result = method() + if isinstance(result, str) and os.path.isfile(result): + return result + except Exception: + continue + + # Common video object attributes that might contain file path + path_attrs = ['filename', 'file_path', 'filepath', 'path', 'source', 'video_path', 'input_path', 'source_file'] + + for attr in path_attrs: + if hasattr(video_data, attr): + try: + value = getattr(video_data, attr) + if isinstance(value, str) and os.path.isfile(value): + return value + except Exception: + continue + + # If it's a dictionary-like object + if hasattr(video_data, '__getitem__'): + for key in path_attrs: + try: + value = video_data[key] + if isinstance(value, str) and os.path.isfile(value): + return value + except (KeyError, TypeError): + continue + + return None + + except Exception: + return None + + def save_video_data(self, video_data, output_path): + """Save video data to file - unified method""" + try: + # First try to extract source file path and copy directly + source_path = self.extract_video_source_path(video_data) + if source_path: + shutil.copy2(source_path, output_path) + return + + # Check if it's a ComfyUI VideoFromFile object with save_to method + if hasattr(video_data, 'save_to'): + try: + # Try to use the save_to method + video_data.save_to(output_path) + return + except Exception: + # Continue to other methods + pass + + # If no source path and no save_to, try to process as tensor data + self.save_video_tensor(video_data, output_path) + + except Exception as e: + print(f"Error in save_video_data: {str(e)}") + raise + + def save_image_tensor(self, image_tensor, output_path): + """Save image tensor to file with animation support""" + try: + from PIL import Image + import numpy as np + import torch + + # Check if we have multiple frames (animation) + has_animation = False + if isinstance(image_tensor, torch.Tensor): + image_np = image_tensor.cpu().numpy() + else: + image_np = image_tensor + + # Determine output format + ext = os.path.splitext(output_path)[1].lower().lstrip('.') + format_map = { + 'jpg': 'JPEG', + 'jpeg': 'JPEG', + 'png': 'PNG', + 'gif': 'GIF', + 'webp': 'WebP', + 'bmp': 'BMP', + 'apng': 'PNG' + } + save_format = format_map.get(ext, 'PNG') + + # Check for animation frames + if len(image_np.shape) == 4 and image_np.shape[0] > 1: + # Multiple frames - handle as animation + frames = [] + for i in range(image_np.shape[0]): + frame_np = image_np[i] + + # Convert to uint8 + if frame_np.dtype == np.float32 or frame_np.dtype == np.float64: + frame_np = (frame_np * 255).astype(np.uint8) + + # Convert to PIL Image + if len(frame_np.shape) == 3 and frame_np.shape[2] == 3: + pil_frame = Image.fromarray(frame_np, 'RGB') + elif len(frame_np.shape) == 3 and frame_np.shape[2] == 4: + pil_frame = Image.fromarray(frame_np, 'RGBA') + elif len(frame_np.shape) == 3 and frame_np.shape[2] == 1: + pil_frame = Image.fromarray(frame_np.squeeze(2), 'L') + else: + pil_frame = Image.fromarray(frame_np) + + frames.append(pil_frame) + + # Save animation + if save_format == 'GIF': + frames[0].save( + output_path, + format='GIF', + save_all=True, + append_images=frames[1:], + duration=100, # 100ms per frame + loop=0 + ) + elif save_format == 'WebP': + frames[0].save( + output_path, + format='WebP', + save_all=True, + append_images=frames[1:], + duration=100, # 100ms per frame + loop=0 + ) + elif ext == 'apng' or save_format == 'PNG': + # APNG support - fallback to first frame if APNG not supported + try: + frames[0].save( + output_path, + format='PNG', + save_all=True, + append_images=frames[1:], + duration=100 + ) + except: + # Fallback to first frame only + frames[0].save(output_path, format='PNG') + print(f"Warning: APNG not supported, saved first frame only") + else: + # Format doesn't support animation, save first frame + pil_image = frames[0] + if save_format == 'JPEG' and pil_image.mode == 'RGBA': + background = Image.new('RGB', pil_image.size, (255, 255, 255)) + background.paste(pil_image, mask=pil_image.split()[-1]) + pil_image = background + pil_image.save(output_path, format=save_format) + print(f"Warning: {save_format} doesn't support animation, saved first frame only") + + else: + # Single frame + if len(image_np.shape) == 4: + image_np = image_np[0] # Take first frame + + # Convert to uint8 + if image_np.dtype == np.float32 or image_np.dtype == np.float64: + image_np = (image_np * 255).astype(np.uint8) + + # Convert to PIL Image + if len(image_np.shape) == 3 and image_np.shape[2] == 3: + pil_image = Image.fromarray(image_np, 'RGB') + elif len(image_np.shape) == 3 and image_np.shape[2] == 4: + pil_image = Image.fromarray(image_np, 'RGBA') + elif len(image_np.shape) == 3 and image_np.shape[2] == 1: + pil_image = Image.fromarray(image_np.squeeze(2), 'L') + elif len(image_np.shape) == 2: + pil_image = Image.fromarray(image_np, 'L') + else: + pil_image = Image.fromarray(image_np) + + # Handle JPEG conversion from RGBA + if save_format == 'JPEG' and pil_image.mode == 'RGBA': + background = Image.new('RGB', pil_image.size, (255, 255, 255)) + background.paste(pil_image, mask=pil_image.split()[-1]) + pil_image = background + + pil_image.save(output_path, format=save_format) + + except ImportError: + print("Warning: PIL not available, saving as pickle") + import pickle + with open(output_path, 'wb') as f: + pickle.dump(image_tensor, f) + + def save_video_tensor(self, video_tensor, output_path): + """Save video tensor to file""" + try: + import cv2 + import numpy as np + import torch + + # Convert tensor to numpy + if isinstance(video_tensor, torch.Tensor): + video_np = video_tensor.cpu().numpy() + else: + video_np = video_tensor + + # Handle different video tensor formats + if len(video_np.shape) == 4: # [frames, height, width, channels] + frames, height, width, channels = video_np.shape + else: + print("Warning: Unexpected video tensor shape, saving as pickle") + import pickle + with open(output_path, 'wb') as f: + pickle.dump(video_tensor, f) + return + + # Convert to uint8 + if video_np.dtype == np.float32 or video_np.dtype == np.float64: + video_np = (video_np * 255).astype(np.uint8) + + # Determine codec from extension + ext = os.path.splitext(output_path)[1].lower().lstrip('.') + if ext == 'mp4': + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + elif ext == 'avi': + fourcc = cv2.VideoWriter_fourcc(*'XVID') + else: + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + + # Create video writer + fps = 30 # Default FPS + out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) + + if not out.isOpened(): + raise RuntimeError(f"Could not open video writer for {output_path}") + + # Write frames + for frame_idx in range(frames): + frame = video_np[frame_idx] + + # Convert RGB to BGR for OpenCV + if channels == 3: + frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) + else: + frame_bgr = frame + + out.write(frame_bgr) + + out.release() + + except ImportError: + print("Warning: OpenCV not available, saving as pickle") + import pickle + with open(output_path, 'wb') as f: + pickle.dump(video_tensor, f) + + diff --git a/requirements.txt b/requirements.txt index 472b405..017376e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,10 @@ # Core dependencies (usually available in ComfyUI) torch>=1.13.0 numpy>=1.21.0 +Pillow>=8.0.0 -# No additional dependencies required -# This package uses only built-in ComfyUI functionality \ No newline at end of file +# Optional dependencies for SaveByFileName node +# For video saving functionality +# opencv-python>=4.5.0 + +# All other nodes use built-in ComfyUI functionality