327 lines
11 KiB
Python
327 lines
11 KiB
Python
|
|
|
|
import os
|
|
import sys
|
|
from .mz_utils import Utils
|
|
from nodes import MAX_RESOLUTION
|
|
import comfy.utils
|
|
import shutil
|
|
import comfy.samplers
|
|
|
|
AUTHOR_NAME = u"[MinusZone]"
|
|
CATEGORY_NAME = f"{AUTHOR_NAME} Utils"
|
|
|
|
sys.path.append(os.path.join(os.path.dirname(__file__)))
|
|
|
|
import importlib
|
|
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
}
|
|
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
}
|
|
|
|
|
|
import mz_llama3
|
|
import mz_llama_cpp
|
|
import mz_llava
|
|
|
|
|
|
|
|
|
|
|
|
llama3_models = [
|
|
"Meta-Llama-3-8B-Instruct.Q4_K_M.gguf",
|
|
"Meta-Llama-3-8B-Instruct.Q2_K.gguf",
|
|
"Meta-Llama-3-8B-Instruct.Q3_K_L.gguf",
|
|
"Meta-Llama-3-8B-Instruct.Q3_K_M.gguf",
|
|
"Meta-Llama-3-8B-Instruct.Q3_K_S.gguf",
|
|
"Meta-Llama-3-8B-Instruct.Q4_0.gguf",
|
|
"Meta-Llama-3-8B-Instruct.Q4_1.gguf",
|
|
"Meta-Llama-3-8B-Instruct.Q4_K_S.gguf",
|
|
"Meta-Llama-3-8B-Instruct.Q5_0.gguf",
|
|
"Meta-Llama-3-8B-Instruct.Q5_1.gguf",
|
|
"Meta-Llama-3-8B-Instruct.Q5_K_M.gguf",
|
|
"Meta-Llama-3-8B-Instruct.Q5_K_S.gguf",
|
|
"Meta-Llama-3-8B-Instruct.Q6_K.gguf",
|
|
"Meta-Llama-3-8B-Instruct.Q8_0.gguf",
|
|
]
|
|
|
|
class MZ_LLama3CLIPTextEncode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
m_llama3_models = llama3_models.copy()
|
|
for i in range(len(m_llama3_models)):
|
|
if mz_llama3.get_exist_model(m_llama3_models[i]) is not None:
|
|
m_llama3_models[i] += "[downloaded]"
|
|
|
|
|
|
importlib.reload(mz_llama3)
|
|
style_presets = mz_llama3.get_style_presets()
|
|
|
|
return {
|
|
"required": {
|
|
"llama3_model": (m_llama3_models, {"default": m_llama3_models[0]}),
|
|
"download_source": (
|
|
["none", "modelscope", "hf-mirror.com",],
|
|
{"default": "none"}
|
|
),
|
|
"style_presets": (
|
|
style_presets, {"default": style_presets[0]}
|
|
),
|
|
"text": ("STRING", {"multiline": True,}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
"optional": {
|
|
"clip": ("CLIP", ),
|
|
},
|
|
}
|
|
RETURN_TYPES = ("STRING", "CONDITIONING",)
|
|
FUNCTION = "encode"
|
|
CATEGORY = CATEGORY_NAME
|
|
def encode(self, text, llama3_model, style_presets, clip=None, download_source=None, seed=0):
|
|
importlib.reload(mz_llama3)
|
|
|
|
llama3_model = llama3_model.replace("[downloaded]", "")
|
|
|
|
text = mz_llama3.query_beautify_prompt_text(llama3_model, text, style_presets, download_source,)
|
|
conditionings = None
|
|
if clip is not None:
|
|
tokens = clip.tokenize(text)
|
|
cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True,)
|
|
conditionings = [[cond, {"pooled_output": pooled}]]
|
|
|
|
return (text, conditionings)
|
|
|
|
NODE_CLASS_MAPPINGS["MZ_LLama3CLIPTextEncode"] = MZ_LLama3CLIPTextEncode
|
|
NODE_DISPLAY_NAME_MAPPINGS["MZ_LLama3CLIPTextEncode"] = f"{AUTHOR_NAME} - LLama3CLIPTextEncode"
|
|
|
|
|
|
LLava_models = [
|
|
"ggml_bakllava-1/ggml-model-q4_k.gguf",
|
|
"ggml_bakllava-1/ggml-model-q5_k.gguf",
|
|
"ggml_bakllava-1/ggml-model-f16.gguf",
|
|
"ggml_llava-v1.5-7b/ggml-model-q4_k.gguf",
|
|
"ggml_llava-v1.5-7b/ggml-model-q5_k.gguf",
|
|
"ggml_llava-v1.5-7b/ggml-model-f16.gguf",
|
|
]
|
|
|
|
LLava_mmproj_models = [
|
|
# "llava-1.6-mistral-7b-gguf/mmproj-model-f16.gguf",
|
|
"ggml_bakllava-1/mmproj-model-f16.gguf",
|
|
"ggml_llava-v1.5-7b/mmproj-model-f16.gguf",
|
|
]
|
|
|
|
class MZ_LLavaImageInterrogator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
|
|
m_llava_models = LLava_models.copy()
|
|
for i in range(len(m_llava_models)):
|
|
if mz_llava.get_exist_model(m_llava_models[i]) is not None:
|
|
m_llava_models[i] += "[downloaded]"
|
|
|
|
m_llava_mmproj_models = LLava_mmproj_models.copy()
|
|
for i in range(len(m_llava_mmproj_models)):
|
|
if mz_llava.get_exist_model(m_llava_mmproj_models[i]) is not None:
|
|
m_llava_mmproj_models[i] += "[downloaded]"
|
|
|
|
|
|
return {"required": {
|
|
"llava_model": (m_llava_models, {"default": m_llava_models[0]}),
|
|
"mmproj_model": (m_llava_mmproj_models, {"default": m_llava_mmproj_models[0]}),
|
|
"download_source": (
|
|
[
|
|
"none",
|
|
"modelscope",
|
|
"hf-mirror.com",
|
|
],
|
|
{"default": "none"}
|
|
),
|
|
"image": ("IMAGE",),
|
|
"resolution": ("INT", {"default": 512, "min": 128, "max": 2048}),
|
|
"prefix" : ("STRING", {"default": "(masterpiece)", "multiline": True, }),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
"optional": {
|
|
"clip": ("CLIP",),
|
|
}}
|
|
RETURN_TYPES = ("STRING", "CONDITIONING",)
|
|
FUNCTION = "interrogate"
|
|
CATEGORY = CATEGORY_NAME
|
|
def interrogate(self, llava_model, mmproj_model, image, resolution, prefix, download_source=None, seed=0, clip=None):
|
|
importlib.reload(mz_llava)
|
|
|
|
llava_model = llava_model.replace("[downloaded]", "")
|
|
mmproj_model = mmproj_model.replace("[downloaded]", "")
|
|
|
|
prefix = Utils.prompt_zh_to_en(prefix)
|
|
if not prefix.endswith(","):
|
|
prefix += ","
|
|
|
|
image_pil = Utils.tensor2pil(image)
|
|
response = mz_llava.image_interrogator(
|
|
llava_model,
|
|
mmproj_model,
|
|
image_pil,
|
|
resolution,
|
|
download_source,
|
|
)
|
|
conditionings = None
|
|
if clip is not None:
|
|
tokens = clip.tokenize(response)
|
|
cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True,)
|
|
conditionings = [[cond, {"pooled_output": pooled}]]
|
|
return (response, conditionings)
|
|
|
|
NODE_CLASS_MAPPINGS["MZ_LLavaImageInterrogator"] = MZ_LLavaImageInterrogator
|
|
NODE_DISPLAY_NAME_MAPPINGS["MZ_LLavaImageInterrogator"] = f"{AUTHOR_NAME} - LLavaImageInterrogator"
|
|
|
|
|
|
qwen_models = [
|
|
"Qwen1.5-1.8B-Chat-GPTQ-Int4",
|
|
"Qwen1.5-7B-Chat-GPTQ-Int4",
|
|
"Qwen1.5-14B-Chat-GPTQ-Int4",
|
|
"Qwen1.5-32B-Chat-GPTQ-Int4",
|
|
"Qwen1.5-72B-Chat-GPTQ-Int4",
|
|
]
|
|
|
|
class MZ_QWenOptionalArgs:
|
|
model_id = ""
|
|
download_source = ""
|
|
customize_prompt = ""
|
|
|
|
class MZ_QWenOptional:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
# qwen_model
|
|
return {
|
|
"required": {
|
|
"model_id": ("STRING", {"default": ""}),
|
|
"customize_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
},
|
|
}
|
|
RETURN_TYPES = ("MZ_QWenOptionalArgs",)
|
|
RETURN_NAMES = ("qwen_optional",)
|
|
FUNCTION = "optional"
|
|
CATEGORY = CATEGORY_NAME
|
|
def optional(self, model_id, download_source, customize_prompt=""):
|
|
result = MZ_QWenOptionalArgs()
|
|
result.model_id = model_id
|
|
result.download_source = download_source
|
|
result.customize_prompt = customize_prompt
|
|
return result
|
|
|
|
# NODE_CLASS_MAPPINGS["MZ_QWenOptional"] = MZ_QWenOptional
|
|
# NODE_DISPLAY_NAME_MAPPINGS["MZ_QWenOptional"] = f"{AUTHOR_NAME} - QWenOptionalParameter"
|
|
|
|
|
|
|
|
class MZ_QWenCLIPTextEncode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"qwen_model": (qwen_models, {"default": "Qwen1.5-1.8B-Chat-GPTQ-Int4"}),
|
|
"download_source": (
|
|
["none", "modelscope", "https://hf-mirror.com",],
|
|
{"default": "none"}
|
|
),
|
|
"text": ("STRING", {"multiline": True,}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
"optional": {
|
|
"clip": ("CLIP", ),
|
|
"qwen_optional": ("MZ_QWenOptionalArgs",),
|
|
},
|
|
}
|
|
RETURN_TYPES = ("STRING", "CONDITIONING",)
|
|
FUNCTION = "encode"
|
|
CATEGORY = CATEGORY_NAME
|
|
def encode(self, clip, text, qwen_model, download_source=None, seed=0, qwen_optional=None):
|
|
importlib.reload(mz_qwen)
|
|
text = mz_qwen.auto_prompt_text(qwen_model, text, download_source, qwen_optional=qwen_optional,)
|
|
conditionings = None
|
|
if clip is not None:
|
|
tokens = clip.tokenize(text)
|
|
cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True,)
|
|
conditionings = [[cond, {"pooled_output": pooled}]]
|
|
|
|
return (text, conditionings)
|
|
|
|
|
|
# NODE_CLASS_MAPPINGS["MZ_QWenCLIPTextEncode"] = MZ_QWenCLIPTextEncode
|
|
# NODE_DISPLAY_NAME_MAPPINGS["MZ_QWenCLIPTextEncode"] = f"{AUTHOR_NAME} - QWenCLIPTextEncode"
|
|
|
|
|
|
|
|
class MZ_QWenImageInterrogator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"image": ("IMAGE",),
|
|
"resolution": ("INT", {"default": 512, "min": 128, "max": 2048}),
|
|
"prefix" : ("STRING", {"default": "(masterpiece)", "multiline": True, }),
|
|
"download_source": (
|
|
["none", "modelscope", "https://hf-mirror.com",],
|
|
{"default": "none"}
|
|
),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
"optional": {
|
|
"clip": ("CLIP",),
|
|
"qwen_optional": ("MZ_QWenOptionalArgs",),
|
|
}}
|
|
RETURN_TYPES = ("STRING", "CONDITIONING",)
|
|
FUNCTION = "interrogate"
|
|
CATEGORY = CATEGORY_NAME
|
|
def interrogate(self, image, resolution, prefix, download_source=None, seed=0, clip=None, qwen_optional=None):
|
|
importlib.reload(mz_qwen)
|
|
prefix = Utils.prompt_zh_to_en(prefix)
|
|
if not prefix.endswith(","):
|
|
prefix += ","
|
|
response = prefix + mz_qwen.auto_image_to_text(image, resolution, download_source, qwen_optional=qwen_optional,)
|
|
|
|
conditionings = None
|
|
if clip is not None:
|
|
tokens = clip.tokenize(response)
|
|
cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
|
|
conditionings = [[cond, {"pooled_output": pooled}]]
|
|
return (response, conditionings)
|
|
|
|
# NODE_CLASS_MAPPINGS["MZ_QWenImageInterrogator"] = MZ_QWenImageInterrogator
|
|
# NODE_DISPLAY_NAME_MAPPINGS["MZ_QWenImageInterrogator"] = f"{AUTHOR_NAME} - QWenImageInterrogator"
|
|
|
|
|
|
|
|
|
|
|
|
class MZ_LLamaCPPInterrogator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_file": ("STRING", {"default": ""}),
|
|
"use_system": ([True, False], {"default": True}),
|
|
"system": ("STRING", {"multiline": True, "default": "You are a helpful assistant."}),
|
|
"prompt": ("STRING", {"multiline": True,}),
|
|
},
|
|
}
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "simple_interrogator"
|
|
CATEGORY = CATEGORY_NAME
|
|
def simple_interrogator(self, model_file, prompt, use_system=True, system="You are a helpful assistant."):
|
|
importlib.reload(mz_llama_cpp)
|
|
result = mz_llama_cpp.llama_cpp_simple_interrogator(
|
|
model_file,
|
|
use_system=use_system,
|
|
system=system,
|
|
question=prompt,
|
|
)
|
|
return (result,)
|
|
|
|
NODE_CLASS_MAPPINGS["MZ_LLamaCPPInterrogator"] = MZ_LLamaCPPInterrogator
|
|
NODE_DISPLAY_NAME_MAPPINGS["MZ_LLamaCPPInterrogator"] = f"{AUTHOR_NAME} - LLamaCPP simple interrogator" |