push

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"
-}
+}