83 lines
2.8 KiB
Python
83 lines
2.8 KiB
Python
import comfy
|
|
import folder_paths
|
|
import nodes
|
|
import aiohttp
|
|
import json
|
|
import asyncio
|
|
from typing import Optional
|
|
from aiohttp import ClientSession, ClientError
|
|
|
|
|
|
class LLM_Loader:
|
|
"""
|
|
Custom node for loading LLM models via a base URL and model name
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"base_url": ([
|
|
"Qwen/通义千问",
|
|
"DeepSeek/深度求索",
|
|
"DouBao/豆包",
|
|
"Spark/星火",
|
|
"GLM/智谱清言",
|
|
"Moonshot/月之暗面",
|
|
"Baichuan/百川",
|
|
"MiniMax/MiniMax",
|
|
"StepFun/阶跃星辰"
|
|
], {}),
|
|
"model": ("STRING", {
|
|
"default": "",
|
|
"label": "模型名称",
|
|
"allow_edit": True
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("base_url", "model")
|
|
FUNCTION = "generate"
|
|
CATEGORY = "DeepSeek_Toolkit"
|
|
|
|
# 定义 base_url 和模型的映射关系
|
|
MODEL_MAPPING = {
|
|
"Qwen/通义千问": ["qwen-turbo", "qwen-plus", "qwen-max"],
|
|
"DeepSeek/深度求索": ["deepseek-7b", "deepseek-33b"],
|
|
"DouBao/豆包": ["doubao-lite", "doubao-pro"],
|
|
"Spark/星火": ["spark-mini", "spark-max"],
|
|
"GLM/智谱清言": ["glm-130b"],
|
|
"Moonshot/月之暗面": ["moonshot-v1"],
|
|
"Baichuan/百川": ["baichuan-7b", "baichuan-13b"],
|
|
"MiniMax/MiniMax": ["minimax-turbo", "minimax-plus"],
|
|
"StepFun/阶跃星辰": ["stepfun-base", "stepfun-large"]
|
|
}
|
|
|
|
def generate(self, base_url: str, model: str):
|
|
# 定义 base_url 映射表
|
|
base_url_mapping = {
|
|
"Qwen/通义千问": "https://dashscope.aliyuncs.com/compatible-mode/v1",
|
|
"DeepSeek/深度求索": "https://api.deepseek.com/v1",
|
|
"DouBao/豆包": "https://ark.cn-beijing.volces.com/api/v3",
|
|
"Spark/星火": "https://spark-api-open.xf-yun.com/v1",
|
|
"GLM/智谱清言": "https://open.bigmodel.cn/api/paas/v4/",
|
|
"Moonshot/月之暗面": "https://api.moonshot.cn/v1",
|
|
"Baichuan/百川": "https://api.baichuan-ai.com/v1",
|
|
"MiniMax/MiniMax": "https://api.minimax.chat/v1",
|
|
"StepFun/阶跃星辰": "https://api.stepfun.com/v1"
|
|
}
|
|
|
|
# 获取实际的 base_url
|
|
actual_base_url = base_url_mapping.get(base_url, base_url)
|
|
|
|
# 返回 base_url 和 model 参数
|
|
return (actual_base_url, model)
|
|
|
|
# 注册节点
|
|
NODE_CLASS_MAPPINGS = {"LLM_Loader": LLM_Loader}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"LLM_Loader": "LLM Loader"}
|
|
|
|
WEB_DIRECTORY = "./web"
|
|
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|