123 lines
4.4 KiB
Python
123 lines
4.4 KiB
Python
import os
|
|
import torch
|
|
import folder_paths
|
|
from typing import Tuple
|
|
from transformers import AutoModelForCausalLM, AutoProcessor # 添加正确导入
|
|
from janus.models import MultiModalityCausalLM, VLChatProcessor # type: ignore
|
|
|
|
class DeepSeekModelLoader:
|
|
"""Model loader node for DeepSeek Janus Pro"""
|
|
|
|
def __init__(self):
|
|
self.local_model_path = None
|
|
|
|
@staticmethod
|
|
def get_available_models():
|
|
"""Scan models/deepseek_janus/Janus-Pro directory and return subfolder names"""
|
|
import os
|
|
models_path = folder_paths.models_dir
|
|
deepseek_janus_path = os.path.join(models_path, "deepseek_janus", "Janus-Pro")
|
|
if not os.path.exists(deepseek_janus_path):
|
|
raise ValueError(
|
|
f"目录未找到: {deepseek_janus_path}. "
|
|
"请确保模型放置在正确的位置。"
|
|
)
|
|
# 获取所有子文件夹名称
|
|
model_names = [
|
|
name for name in os.listdir(deepseek_janus_path)
|
|
if os.path.isdir(os.path.join(deepseek_janus_path, name))
|
|
]
|
|
if not model_names:
|
|
raise ValueError(
|
|
f"{deepseek_janus_path} 中未找到模型. "
|
|
"请下载并将模型放置在正确的位置。"
|
|
)
|
|
return model_names
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_name": (cls.get_available_models(),),
|
|
"use_local": ([True, False], {
|
|
"default": True
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("DEEPSEEK_MODEL", "DEEPSEEK_TOKENIZER")
|
|
RETURN_NAMES = ("model", "tokenizer")
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "DeepSeek_Toolkit/Loader"
|
|
|
|
def load_model(self, model_name: str, use_local: bool) -> Tuple:
|
|
"""Load DeepSeek Janus Pro model and processor"""
|
|
# 设置设备和数据类型
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
try:
|
|
dtype = torch.bfloat16
|
|
torch.zeros(1, dtype=dtype, device=device)
|
|
except RuntimeError:
|
|
dtype = torch.float16
|
|
|
|
try:
|
|
if use_local:
|
|
# 获取ComfyUI根目录
|
|
models_path = folder_paths.models_dir
|
|
# 构建模型路径
|
|
model_dir = os.path.join(models_path, "deepseek_janus", "Janus-Pro", model_name)
|
|
|
|
if not os.path.exists(model_dir):
|
|
raise ValueError(
|
|
f"本地模型未在 {model_dir} 找到. "
|
|
"请下载模型并将其放置在 ComfyUI/models/deepseek_janus/Janus-Pro 文件夹中。"
|
|
)
|
|
model_path = model_dir
|
|
else:
|
|
model_path = model_name
|
|
|
|
print(f"从 {model_path} 加载 DeepSeek Janus 模型")
|
|
|
|
# 加载处理器和模型
|
|
try:
|
|
processor = VLChatProcessor.from_pretrained(
|
|
str(model_path),
|
|
trust_remote_code=True,
|
|
use_fast=False
|
|
)
|
|
except Exception as e:
|
|
raise ValueError(
|
|
f"无法从 {model_path} 加载处理器. 请确保已下载正确的模型文件并将其放置在 models/deepseek_janus/Janus-Pro 文件夹中. 错误：{str(e)}"
|
|
)
|
|
|
|
try:
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
str(model_path),
|
|
trust_remote_code=True,
|
|
torch_dtype=dtype,
|
|
device_map="auto"
|
|
)
|
|
except Exception as e:
|
|
raise ValueError(
|
|
f"无法从 {model_path} 加载模型. 请确保已下载正确的模型文件并将其放置在 models/deepseek_janus/Janus-Pro 文件夹中. 错误：{str(e)}"
|
|
)
|
|
|
|
# 将模型移动到正确的设备和数据类型
|
|
model = model.to(dtype).to(device).eval()
|
|
|
|
print("DeepSeek Janus 模型加载成功")
|
|
return (model, processor)
|
|
|
|
except Exception as e:
|
|
print(f"加载模型时出错: {str(e)}")
|
|
raise
|
|
|
|
# 节点类映射
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DeepSeekModelLoader": DeepSeekModelLoader
|
|
}
|
|
|
|
# 节点显示名称映射
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DeepSeekModelLoader": "Multimodal Loader"
|
|
} |