266 lines
8.0 KiB
Python
266 lines
8.0 KiB
Python
import os
|
|
import random
|
|
import string
|
|
from dashscope import Generation, MultiModalConversation
|
|
from http import HTTPStatus
|
|
import torchvision.transforms as transforms
|
|
|
|
|
|
def generate_random_image_name():
|
|
random_number = "".join(random.choices(string.ascii_letters + string.digits, k=10))
|
|
return f"image_{random_number}.png"
|
|
|
|
|
|
def get_image_url(image):
|
|
image = image.squeeze(0)
|
|
image = image.permute(2, 0, 1)
|
|
image = transforms.ToPILImage()(image)
|
|
|
|
image_name = generate_random_image_name()
|
|
image.save(image_name)
|
|
return image_name
|
|
|
|
|
|
def get_model_versions(model_type: str) -> list[str]:
|
|
model_versions = []
|
|
current_group = None
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
model_versions_file = os.path.join(
|
|
current_dir, "model_versions", f"{model_type}.txt"
|
|
)
|
|
|
|
with open(model_versions_file, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line or line.startswith("#"):
|
|
continue
|
|
|
|
if line.startswith("[") and line.endswith("]"):
|
|
current_group = line[1:-1]
|
|
continue
|
|
|
|
model_versions.append(line)
|
|
|
|
return model_versions
|
|
|
|
|
|
class DashscopeLLMLoader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
model_versions = get_model_versions("llm")
|
|
return {
|
|
"required": {
|
|
"model_version": (
|
|
model_versions,
|
|
{"default": model_versions[0]},
|
|
),
|
|
}
|
|
}
|
|
|
|
FUNCTION = "select_model"
|
|
OUTPUT_NODE = True
|
|
RETURN_TYPES = ("STRING",)
|
|
|
|
CATEGORY = "dashscope"
|
|
|
|
def select_model(self, model_version) -> tuple[str]:
|
|
if not model_version:
|
|
raise ValueError("Model version cannot be empty")
|
|
return (model_version,)
|
|
|
|
|
|
class DashscopeVLMLoader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
model_versions = get_model_versions("vlm")
|
|
return {
|
|
"required": {
|
|
"model_version": (
|
|
model_versions,
|
|
{"default": model_versions[0]},
|
|
),
|
|
}
|
|
}
|
|
|
|
FUNCTION = "select_model"
|
|
OUTPUT_NODE = True
|
|
RETURN_TYPES = ("STRING",)
|
|
|
|
CATEGORY = "dashscope"
|
|
|
|
def select_model(self, model_version) -> tuple[str]:
|
|
if not model_version:
|
|
raise ValueError("Model version cannot be empty")
|
|
return (model_version,)
|
|
|
|
|
|
# 关于OCR模型的最小像素值，根据阿里云文字识别服务的要求，图片的最短边不小于15像素，单字大小在10-50像素内时，识别效果较好。因此，28像素并不是OCR模型的最小像素值。
|
|
# 您好，关于OCR中min_pixels为3136（即28284）的28、28和4的含义如下：
|
|
|
|
# 28：表示图像的高度，单位为像素。
|
|
# 28：表示图像的宽度，单位为像素。
|
|
# 4：表示每个像素的通道数，通常为4个通道（R、G、B和Alpha）。
|
|
|
|
|
|
# 关于OCR中图片的宽高比要求，根据阿里云文字识别服务的要求，单张图片的最长边不超过4096像素，最短边不小于15像素，当长边超过1024像素时，长宽比不超过1:10。
|
|
|
|
|
|
# 这通常与OCR模型的设计和性能优化有关。1280这个值可能是为了平衡图像质量和处理速度，确保在大多数情况下能够高效地处理图像。
|
|
class DashscopeOCRCaller:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
model_versions = get_model_versions("ocr")
|
|
return {
|
|
"required": {
|
|
"model_version": (
|
|
model_versions,
|
|
{"default": model_versions[0]},
|
|
),
|
|
"min_pixels": (
|
|
"STRING",
|
|
{"default": 28 * 28 * 4},
|
|
), # width, height, channels
|
|
"max_pixels": (
|
|
"STRING",
|
|
{"default": 28 * 28 * 1280},
|
|
), # larger channels to balance image quality and processing speed
|
|
"image": ("IMAGE", {"default": None}),
|
|
}
|
|
}
|
|
|
|
FUNCTION = "call_model"
|
|
OUTPUT_NODE = True
|
|
RETURN_TYPES = ("STRING",)
|
|
|
|
CATEGORY = "dashscope"
|
|
|
|
def call_model(self, model_version, min_pixels, max_pixels, image) -> tuple[str]:
|
|
if not model_version:
|
|
raise ValueError("Model version cannot be empty")
|
|
|
|
api_key = os.getenv("DASHSCOPE_API_KEY")
|
|
if not api_key:
|
|
raise ValueError("DASHSCOPE_API_KEY environment variable is not set")
|
|
|
|
print("Call the OCR model")
|
|
image_url = get_image_url(image)
|
|
messages = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"image": image_url,
|
|
"min_pixels": min_pixels,
|
|
"max_pixels": max_pixels,
|
|
},
|
|
{"type": "text", "text": "Read all the text in the image."},
|
|
],
|
|
},
|
|
]
|
|
response = MultiModalConversation.call(
|
|
api_key=api_key,
|
|
model=model_version,
|
|
messages=messages,
|
|
result_format="message",
|
|
)
|
|
|
|
if response.status_code == HTTPStatus.OK:
|
|
print(response.output)
|
|
else:
|
|
print(response.code)
|
|
print(response.message)
|
|
|
|
message_content = response.output.choices[0].message.content[0]["text"]
|
|
return (message_content,)
|
|
|
|
|
|
class DashscopeModelCaller:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_version": ("STRING", {"default": None}),
|
|
"system_prompt": (
|
|
"STRING",
|
|
{"default": None},
|
|
),
|
|
"user_prompt": ("STRING", {"default": None}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE", {"default": None}),
|
|
},
|
|
}
|
|
|
|
FUNCTION = "call_model"
|
|
OUTPUT_NODE = True
|
|
RETURN_TYPES = ("STRING",)
|
|
|
|
CATEGORY = "dashscope"
|
|
|
|
def call_model(
|
|
self, model_version, system_prompt, user_prompt, image
|
|
) -> tuple[str]:
|
|
if not model_version:
|
|
raise ValueError("Model version cannot be empty")
|
|
if not system_prompt:
|
|
raise ValueError("System prompt cannot be empty")
|
|
if not user_prompt:
|
|
raise ValueError("User prompt cannot be empty")
|
|
|
|
api_key = os.getenv("DASHSCOPE_API_KEY")
|
|
if not api_key:
|
|
raise ValueError("DASHSCOPE_API_KEY environment variable is not set")
|
|
|
|
if image is None:
|
|
print("Call the LLM model")
|
|
messages = [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_prompt},
|
|
]
|
|
response = Generation.call(
|
|
api_key=api_key,
|
|
model=model_version,
|
|
messages=messages,
|
|
result_format="message",
|
|
)
|
|
else:
|
|
print("Call the VLM model")
|
|
image_url = get_image_url(image)
|
|
messages = [
|
|
{"role": "system", "content": [{"text": system_prompt}]},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"image": image_url},
|
|
{"text": user_prompt},
|
|
],
|
|
},
|
|
]
|
|
response = MultiModalConversation.call(
|
|
api_key=api_key,
|
|
model=model_version,
|
|
messages=messages,
|
|
result_format="message",
|
|
)
|
|
|
|
if response.status_code == HTTPStatus.OK:
|
|
print(response.output)
|
|
else:
|
|
print(response.code)
|
|
print(response.message)
|
|
|
|
message_content = response.output.choices[0].message.content[0]["text"]
|
|
return (message_content,)
|