youxirenshengandClaude
99e219c55e
- Add ModelMetadataReader node for reading model metadata - Add ModelListReader node for listing models - Support multiple model types (loras, checkpoints, vae, etc.) - Extract metadata from safetensors files including download URLs - Add comprehensive README documentation 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
163 lines
6.2 KiB
Python
163 lines
6.2 KiB
Python
"""
|
|
Model Metadata Reader Nodes
|
|
读取模型文件的元数据信息，包括下载地址等
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import safetensors.torch
|
|
import folder_paths
|
|
|
|
|
|
class ModelMetadataReader:
|
|
"""读取模型文件的元数据信息"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
model_types = ["checkpoints", "loras", "vae", "controlnet", "embeddings",
|
|
"upscale_models", "clip_vision", "diffusion_models", "text_encoders"]
|
|
|
|
return {
|
|
"required": {
|
|
"model_type": (model_types, {"default": "loras"}),
|
|
"model_name": ("STRING", {"default": "", "multiline": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("metadata_json", "download_url", "model_info")
|
|
FUNCTION = "read_metadata"
|
|
CATEGORY = "utils/model"
|
|
DESCRIPTION = "读取模型文件的元数据信息，包括下载地址等"
|
|
|
|
def read_metadata(self, model_type, model_name):
|
|
"""读取模型元数据"""
|
|
metadata_json = ""
|
|
download_url = ""
|
|
model_info = ""
|
|
|
|
try:
|
|
# 获取模型完整路径
|
|
if not model_name:
|
|
model_info = f"请输入模型名称\n可用模型类型: {model_type}"
|
|
return (metadata_json, download_url, model_info)
|
|
|
|
full_path = folder_paths.get_full_path(model_type, model_name)
|
|
|
|
if full_path is None:
|
|
# 尝试列出可用的模型文件
|
|
available_models = folder_paths.get_filename_list(model_type)
|
|
model_info = f"未找到模型: {model_name}\n\n可用模型 ({model_type}):\n"
|
|
model_info += "\n".join(available_models[:20]) # 只显示前20个
|
|
if len(available_models) > 20:
|
|
model_info += f"\n... 还有 {len(available_models) - 20} 个模型"
|
|
return (metadata_json, download_url, model_info)
|
|
|
|
# 读取模型元数据
|
|
metadata = None
|
|
file_ext = os.path.splitext(full_path)[1].lower()
|
|
|
|
if file_ext in ['.safetensors', '.sft']:
|
|
# 读取safetensors文件的元数据
|
|
with safetensors.safe_open(full_path, framework="pt", device="cpu") as f:
|
|
metadata = f.metadata()
|
|
|
|
if metadata:
|
|
# 转换为JSON字符串
|
|
metadata_json = json.dumps(metadata, indent=2, ensure_ascii=False)
|
|
|
|
# 提取下载地址
|
|
download_url = metadata.get("modelspec.download_url",
|
|
metadata.get("ss_download_url",
|
|
metadata.get("download_url", "")))
|
|
|
|
# 生成模型信息摘要
|
|
model_info = f"模型文件: {model_name}\n"
|
|
model_info += f"完整路径: {full_path}\n"
|
|
model_info += f"文件类型: {file_ext}\n"
|
|
model_info += f"文件大小: {os.path.getsize(full_path) / (1024*1024):.2f} MB\n\n"
|
|
|
|
# 添加一些关键元数据
|
|
if "modelspec.title" in metadata:
|
|
model_info += f"标题: {metadata['modelspec.title']}\n"
|
|
if "modelspec.description" in metadata:
|
|
model_info += f"描述: {metadata['modelspec.description']}\n"
|
|
if "modelspec.architecture" in metadata:
|
|
model_info += f"架构: {metadata['modelspec.architecture']}\n"
|
|
if "ss_base_model_version" in metadata:
|
|
model_info += f"基础模型: {metadata['ss_base_model_version']}\n"
|
|
if "ss_resolution" in metadata:
|
|
model_info += f"训练分辨率: {metadata['ss_resolution']}\n"
|
|
if "ss_network_dim" in metadata:
|
|
model_info += f"网络维度: {metadata['ss_network_dim']}\n"
|
|
if "ss_network_alpha" in metadata:
|
|
model_info += f"网络Alpha: {metadata['ss_network_alpha']}\n"
|
|
if download_url:
|
|
model_info += f"\n下载地址: {download_url}\n"
|
|
|
|
# 显示元数据字段数量
|
|
model_info += f"\n元数据字段总数: {len(metadata)}"
|
|
else:
|
|
model_info = f"模型文件: {model_name}\n"
|
|
model_info += f"完整路径: {full_path}\n"
|
|
model_info += f"文件类型: {file_ext}\n"
|
|
model_info += f"文件大小: {os.path.getsize(full_path) / (1024*1024):.2f} MB\n\n"
|
|
model_info += "该文件格式不支持读取元数据或没有元数据信息"
|
|
|
|
except Exception as e:
|
|
model_info = f"读取模型元数据时发生错误:\n{str(e)}"
|
|
|
|
return (metadata_json, download_url, model_info)
|
|
|
|
|
|
class ModelListReader:
|
|
"""列出指定类型的所有模型"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
model_types = ["checkpoints", "loras", "vae", "controlnet", "embeddings",
|
|
"upscale_models", "clip_vision", "diffusion_models", "text_encoders"]
|
|
|
|
return {
|
|
"required": {
|
|
"model_type": (model_types, {"default": "loras"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "INT")
|
|
RETURN_NAMES = ("model_list", "model_count")
|
|
FUNCTION = "list_models"
|
|
CATEGORY = "utils/model"
|
|
DESCRIPTION = "列出指定类型的所有模型文件"
|
|
|
|
def list_models(self, model_type):
|
|
"""列出模型"""
|
|
try:
|
|
# 获取模型列表
|
|
models = folder_paths.get_filename_list(model_type)
|
|
model_count = len(models)
|
|
|
|
# 生成列表字符串
|
|
model_list = f"模型类型: {model_type}\n"
|
|
model_list += f"模型总数: {model_count}\n\n"
|
|
model_list += "模型列表:\n"
|
|
model_list += "\n".join([f"{i+1}. {model}" for i, model in enumerate(models)])
|
|
|
|
except Exception as e:
|
|
model_list = f"读取模型列表时发生错误:\n{str(e)}"
|
|
model_count = 0
|
|
|
|
return (model_list, model_count)
|
|
|
|
|
|
# 节点映射
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ModelMetadataReader": ModelMetadataReader,
|
|
"ModelListReader": ModelListReader,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ModelMetadataReader": "Model Metadata Reader",
|
|
"ModelListReader": "Model List Reader",
|
|
}
|