diff --git a/__pycache__/model_loader.cpython-312.pyc b/__pycache__/model_loader.cpython-312.pyc index 4d93b08..3225aff 100644 Binary files a/__pycache__/model_loader.cpython-312.pyc and b/__pycache__/model_loader.cpython-312.pyc differ diff --git a/model_loader.py b/model_loader.py index a26141f..9d2e95a 100644 --- a/model_loader.py +++ b/model_loader.py @@ -1,20 +1,19 @@ import os import torch -import folder_paths from typing import Tuple -from transformers import AutoModelForCausalLM, AutoProcessor # 添加正确导入 -from janus.models import MultiModalityCausalLM, VLChatProcessor # type: ignore +from transformers import AutoModelForCausalLM, AutoTokenizer +from janus.models import VLChatProcessor # type: ignore +import folder_paths 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): @@ -22,7 +21,6 @@ class DeepSeekModelLoader: 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)) @@ -33,7 +31,7 @@ class DeepSeekModelLoader: "请下载并将模型放置在正确的位置。" ) return model_names - + @classmethod def INPUT_TYPES(cls): return { @@ -44,7 +42,7 @@ class DeepSeekModelLoader: }), }, } - + RETURN_TYPES = ("DEEPSEEK_MODEL", "DEEPSEEK_TOKENIZER") RETURN_NAMES = ("model", "tokenizer") FUNCTION = "load_model" @@ -52,7 +50,6 @@ class DeepSeekModelLoader: 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 @@ -62,11 +59,9 @@ class DeepSeekModelLoader: try: if use_local: - # 获取ComfyUI根目录 models_path = folder_paths.models_dir - # 构建模型路径 - model_dir = os.path.join(models_path, "deepseek_janus", "Janus-Pro", model_name) - + model_dir = os.path.join(models_path, "deepseek_janus", "Janus-Pro", os.path.basename(model_name)) + if not os.path.exists(model_dir): raise ValueError( f"本地模型未在 {model_dir} 找到. " @@ -76,41 +71,24 @@ class DeepSeekModelLoader: else: model_path = model_name + # 确保 model_path 是字符串 + model_path = str(model_path) + 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)}" - ) + + # 加载 tokenizer 和 processor + tokenizer = AutoTokenizer.from_pretrained(model_path) + processor = VLChatProcessor.from_pretrained(model_path, image_processor=None, tokenizer=tokenizer) + + # 加载模型 + model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, torch_dtype=dtype, device_map="auto") - 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 + raise ValueError(f"加载模型时出错: {str(e)}") # 节点类映射 NODE_CLASS_MAPPINGS = { @@ -120,4 +98,4 @@ NODE_CLASS_MAPPINGS = { # 节点显示名称映射 NODE_DISPLAY_NAME_MAPPINGS = { "DeepSeekModelLoader": "Multimodal Loader" -} \ No newline at end of file +}