o1keyandClaude Sonnet 4.5
9ab209b2b7
Complete rewrite/sync of comfyui_o1key custom nodes. Treat this commit as the current canonical version. Co-Authored-By: Claude Sonnet 4.5 <[email protected]>
147 lines
4.4 KiB
Python
147 lines
4.4 KiB
Python
"""
|
||
LoadFile 节点
|
||
ComfyUI 自定义节点,用于加载文件并转换为 FILE 类型数据
|
||
"""
|
||
|
||
import base64
|
||
import os
|
||
from pathlib import Path
|
||
from typing import Tuple
|
||
|
||
from ..utils.file_types import FileData, DOCUMENT_MIME_TYPES, FILE_SIZE_LIMITS
|
||
|
||
|
||
class LoadFile:
|
||
"""
|
||
LoadFile 节点
|
||
|
||
功能：
|
||
- 从文件系统加载文件
|
||
- 支持 PDF 和 TXT 文件
|
||
- 转换为 FILE 类型数据（包含 base64 编码内容）
|
||
- 验证文件大小和格式
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
"""
|
||
定义输入参数
|
||
"""
|
||
return {
|
||
"required": {
|
||
"文件路径": ("STRING", {
|
||
"default": "",
|
||
"multiline": False
|
||
})
|
||
}
|
||
}
|
||
|
||
# 返回值类型
|
||
RETURN_TYPES = ("FILE", "STRING")
|
||
RETURN_NAMES = ("文件", "文件信息")
|
||
|
||
# 执行函数名
|
||
FUNCTION = "load_file"
|
||
|
||
# 节点分类
|
||
CATEGORY = "file/input"
|
||
|
||
def load_file(self, 文件路径: str) -> Tuple[FileData, str]:
|
||
"""
|
||
加载文件并转换为 FILE 类型
|
||
|
||
Args:
|
||
文件路径: 文件的完整路径（支持绝对路径和相对路径）
|
||
|
||
Returns:
|
||
(FileData, 文件信息预览)
|
||
|
||
Raises:
|
||
ValueError: 文件不存在、不支持的文件类型或文件过大
|
||
"""
|
||
try:
|
||
# 清理路径（去除空格和引号）
|
||
file_path = 文件路径.strip().strip('"').strip("'")
|
||
|
||
if not file_path:
|
||
raise ValueError("文件路径不能为空")
|
||
|
||
# 转换为 Path 对象
|
||
path = Path(file_path)
|
||
|
||
# 如果是相对路径，转换为绝对路径
|
||
if not path.is_absolute():
|
||
# 相对于当前工作目录
|
||
path = Path.cwd() / path
|
||
|
||
# 验证文件是否存在
|
||
if not path.exists():
|
||
raise ValueError(f"文件不存在: {file_path}")
|
||
|
||
if not path.is_file():
|
||
raise ValueError(f"路径不是文件: {file_path}")
|
||
|
||
# 获取文件信息
|
||
extension = path.suffix.lower()
|
||
filename = path.stem
|
||
file_size = path.stat().st_size
|
||
|
||
# 验证文件类型
|
||
if extension not in DOCUMENT_MIME_TYPES:
|
||
supported_types = ", ".join(DOCUMENT_MIME_TYPES.keys())
|
||
raise ValueError(
|
||
f"不支持的文件类型: {extension}\n"
|
||
f"支持的类型: {supported_types}"
|
||
)
|
||
|
||
# 获取 MIME 类型
|
||
mime_type = DOCUMENT_MIME_TYPES[extension]
|
||
|
||
# 验证文件大小
|
||
size_limit = FILE_SIZE_LIMITS.get(extension, 20 * 1024 * 1024)
|
||
if file_size > size_limit:
|
||
raise ValueError(
|
||
f"文件过大 ({file_size / 1024 / 1024:.2f}MB)，"
|
||
f"最大支持 {size_limit / 1024 / 1024:.0f}MB"
|
||
)
|
||
|
||
# 读取文件并转换为 base64
|
||
print(f"LoadFile: 正在加载文件 {filename}{extension}")
|
||
print(f"LoadFile: 文件大小 = {file_size / 1024:.2f}KB")
|
||
|
||
with open(path, "rb") as f:
|
||
file_bytes = f.read()
|
||
|
||
# Base64 编码
|
||
b64_str = base64.b64encode(file_bytes).decode("utf-8")
|
||
|
||
# 创建 FileData 对象
|
||
file_data = FileData(
|
||
path=str(path),
|
||
filename=filename,
|
||
extension=extension,
|
||
mime_type=mime_type,
|
||
data=b64_str,
|
||
size=file_size
|
||
)
|
||
|
||
# 生成文件信息预览
|
||
file_info = (
|
||
f"文件名: {filename}{extension}\n"
|
||
f"类型: {mime_type}\n"
|
||
f"大小: {file_size / 1024:.2f}KB\n"
|
||
f"路径: {path}"
|
||
)
|
||
|
||
print(f"LoadFile: 加载成功")
|
||
|
||
return (file_data, file_info)
|
||
|
||
except ValueError as e:
|
||
print(f"LoadFile: 输入错误 - {str(e)}")
|
||
raise
|
||
|
||
except Exception as e:
|
||
print(f"LoadFile: 未知错误 - {str(e)}")
|
||
raise
|