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
This commit is contained in:
@@ -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
|
||||
|
||||
+6
-1
@@ -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 文本集成
|
||||
|
||||
+6
-1
@@ -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 文本整合
|
||||
|
||||
@@ -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)
|
||||
@@ -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__ = [
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
+6
-2
@@ -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
|
||||
# Optional dependencies for SaveByFileName node
|
||||
# For video saving functionality
|
||||
# opencv-python>=4.5.0
|
||||
|
||||
# All other nodes use built-in ComfyUI functionality
|
||||
|
||||
Reference in New Issue
Block a user