nodes/MinusZoneAI-ComfyUI-Prompt-MZ
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

'update'

Browse Source
This commit is contained in:
wailovet
2024-06-24 00:11:38 +08:00
parent 7324dadb1a
commit 06fc84c2ed
5 changed files with 329 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+111 -4
View File
@@ -352,6 +352,7 @@ class MZ_ImageCaptionerConfig:
"caption_suffix": ("STRING", {"default": ".caption"}),
"force_update": ([False, True], {"default": False}),
"retry_keyword": ("STRING", {"default": "not,\",error"}),
"prompt_fixed_beginning": ("STRING", {"default": "", }),
},
"optional": {
@@ -606,16 +607,17 @@ class MZ_Florence2CLIPTextEncode:
def INPUT_TYPES(s):
return {
"required": {
# "model_name": ("STRING", {"default": "Florence-2-large", "placeholder": "model_path"}),
"model_name": ([
"Florence-2-large-ft",
"Florence-2-large",
],),
"resolution": ("INT", {"default": 512, "min": 128, "max": 0xffffffffffffffff}),
"keep_device": ([False, True], {"default": False}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
# "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
},
"optional": {
"image": ("IMAGE",),
"clip": ("CLIP", ),
# "customize_instruct": ("CustomizeInstruct", ),
"captioner_config": ("ImageCaptionerConfig", ),
},
}
@@ -638,6 +640,111 @@ NODE_DISPLAY_NAME_MAPPINGS[
"MZ_Florence2CLIPTextEncode"] = f"{AUTHOR_NAME} - CLIPTextEncode(Florence-2)"
class MZ_Florence2Captioner:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model_name": ([
"Florence-2-large-ft",
"Florence-2-large",
],),
"resolution": ("INT", {"default": 512, "min": 128, "max": 0xffffffffffffffff}),
"captioner_config": ("ImageCaptionerConfig", ),
},
"optional": {
},
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("debug",)
OUTPUT_NODE = True
FUNCTION = "encode"
CATEGORY = CATEGORY_NAME
def encode(self, **kwargs):
kwargs = kwargs.copy()
from . import mz_transformers
importlib.reload(mz_transformers)
return mz_transformers.florence2_node_encode(kwargs)
NODE_CLASS_MAPPINGS["MZ_Florence2Captioner"] = MZ_Florence2Captioner
NODE_DISPLAY_NAME_MAPPINGS[
"MZ_Florence2Captioner"] = f"{AUTHOR_NAME} - Captioner(Florence-2)"
class MZ_PaliGemmaCLIPTextEncode:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model_name": ([
"paligemma-sd3-long-captioner-v2",
"paligemma-sd3-long-captioner",
],),
"resolution": ("INT", {"default": 512, "min": 128, "max": 0xffffffffffffffff}),
"keep_device": ([False, True], {"default": False}),
# "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
},
"optional": {
"image": ("IMAGE",),
"clip": ("CLIP", ),
},
}
RETURN_TYPES = ("STRING", "CONDITIONING",)
RETURN_NAMES = ("text", "conditioning",)
OUTPUT_NODE = True
FUNCTION = "encode"
CATEGORY = CATEGORY_NAME
def encode(self, **kwargs):
kwargs = kwargs.copy()
from . import mz_transformers
importlib.reload(mz_transformers)
return mz_transformers.paligemma_node_encode(kwargs)
NODE_CLASS_MAPPINGS["MZ_PaliGemmaCLIPTextEncode"] = MZ_PaliGemmaCLIPTextEncode
NODE_DISPLAY_NAME_MAPPINGS[
"MZ_PaliGemmaCLIPTextEncode"] = f"{AUTHOR_NAME} - CLIPTextEncode(PaliGemma)"
class MZ_PaliGemmaCaptioner:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model_name": ([
"paligemma-sd3-long-captioner-v2",
"paligemma-sd3-long-captioner",
],),
"resolution": ("INT", {"default": 512, "min": 128, "max": 0xffffffffffffffff}),
"captioner_config": ("ImageCaptionerConfig", ),
},
"optional": {
},
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("debug",)
OUTPUT_NODE = True
FUNCTION = "encode"
CATEGORY = CATEGORY_NAME
def encode(self, **kwargs):
kwargs = kwargs.copy()
from . import mz_transformers
importlib.reload(mz_transformers)
return mz_transformers.paligemma_node_encode(kwargs)
NODE_CLASS_MAPPINGS["MZ_PaliGemmaCaptioner"] = MZ_PaliGemmaCaptioner
NODE_DISPLAY_NAME_MAPPINGS[
"MZ_PaliGemmaCaptioner"] = f"{AUTHOR_NAME} - Captioner(PaliGemma)"
try:
from . import mz_gen_translate
mz_gen_translate.gen_translate(
mz_llama_core_nodes.py
+3 -1
View File
@@ -319,7 +319,9 @@ def image_interrogator_captioner(args_dict):
if response != "":
with open(caption_file, "w") as f:
f.write(response)
prompt_fixed_beginning = captioner_config.get(
"prompt_fixed_beginning", "")
f.write(prompt_fixed_beginning + response)
result.append(response)
mz_prompt_utils.py
+12
View File
@@ -422,6 +422,13 @@ class Utils:
def get_device():
return comfy.model_management.get_torch_device()
def download_small_file(url, filepath):
response = requests.get(url)
os.makedirs(os.path.dirname(filepath), exist_ok=True)
with open(filepath, "wb") as f:
f.write(response.content)
return filepath
def download_file(url, filepath, threads=8, retries=6):
get_size_tmp = requests.get(url, stream=True)
@@ -429,6 +436,11 @@ class Utils:
print(f"Downloading {url} to {filepath} with size {total_size} bytes")
# 如果文件大小小于 50MB,使用download_small_file
if total_size < 50 * 1024 * 1024:
return Utils.download_small_file(url, filepath)
base_filename = os.path.basename(filepath)
cache_dir = os.path.join(os.path.dirname(
filepath), f"{base_filename}.t_{threads}_cache")
mz_transformers.py
+202 -17
View File
@@ -3,7 +3,7 @@ import shutil
import subprocess
def florence2_captioner(args_dict):
def transformers_captioner(args_dict, myfunc):
from . import mz_prompt_utils
import PIL.Image as Image
captioner_config = args_dict.get("captioner_config", {})
@@ -59,13 +59,15 @@ def florence2_captioner(args_dict):
thumbnail,
)
response = florence2_node_encode(onec_args_dict)
response = myfunc(onec_args_dict)
response = response.get("result", ())[0]
response = response.strip()
if response != "":
with open(caption_file, "w") as f:
f.write(response)
prompt_fixed_beginning = captioner_config.get(
"prompt_fixed_beginning", "")
f.write(prompt_fixed_beginning + response)
result.append(response)
return result
@@ -75,7 +77,7 @@ def florence2_node_encode(args_dict):
args_dict = args_dict.copy()
captioner_config = args_dict.get("captioner_config", None)
if captioner_config is not None:
florence2_captioner(args_dict)
transformers_captioner(args_dict, florence2_node_encode)
# raise Exception(
# "图片批量反推任务已完成 ; Image batch reverse push task completed")
return {"ui": {"string": ["图片批量反推任务已完成 ; Image batch reverse push task completed",]}, "result": ("", None)}
@@ -85,24 +87,30 @@ def florence2_node_encode(args_dict):
from . import mz_prompt_utils
from .mz_prompt_utils import Utils
florence2_large_files = [
{
"file_path": "pytorch_model.bin",
"url": "https://www.modelscope.cn/api/v1/models/AI-ModelScope/Florence-2-large/repo?Revision=master&FilePath=pytorch_model.bin"
}
]
florence2_large_files_map = {
"Florence-2-large": [
{
"file_path": "pytorch_model.bin",
"url": "https://www.modelscope.cn/api/v1/models/AI-ModelScope/Florence-2-large/repo?Revision=master&FilePath=pytorch_model.bin"
}
],
"Florence-2-large-ft": [
{
"file_path": "pytorch_model.bin",
"url": "https://www.modelscope.cn/api/v1/models/AI-ModelScope/Florence-2-large-ft/repo?Revision=master&FilePath=pytorch_model.bin"
}
],
}
# model_name = args_dict.get("model_name")
# 判断是否是绝对路径
# if os.path.isabs(model_name):
# model_path = model_name
# else:
llm_path = os.path.join(
folder_paths.models_dir,
"LLM",
)
os.makedirs(llm_path, exist_ok=True)
model_path = os.path.join(llm_path, "Florence-2-large")
model_name = args_dict.get("model_name", "Florence-2-large")
model_path = os.path.join(llm_path, model_name)
if not os.path.exists(model_path):
# GIT_LFS_SKIP_SMUDGE=1 git clone https://www.modelscope.cn/AI-ModelScope/Florence-2-large.git
@@ -113,6 +121,7 @@ def florence2_node_encode(args_dict):
if original_env is not None:
os.environ["GIT_LFS_SKIP_SMUDGE"] = original_env
florence2_large_files = florence2_large_files_map.get(model_name, [])
for file_info in florence2_large_files:
file_path = os.path.join(model_path, file_info["file_path"])
# 判断文件大小小于1M
@@ -136,7 +145,7 @@ def florence2_node_encode(args_dict):
local_files_only=True,
trust_remote_code=True,
)
model.to(device)
model.to(device).eval()
Utils.cache_set(f"florence_model_and_opt_", model)
processor = AutoProcessor.from_pretrained(
@@ -186,3 +195,179 @@ def florence2_node_encode(args_dict):
clip, response, )
return {"ui": {"string": [mz_prompt_utils.Utils.to_debug_prompt(response),]}, "result": (response, conditionings)}
def paligemma_node_encode(args_dict):
args_dict = args_dict.copy()
captioner_config = args_dict.get("captioner_config", None)
if captioner_config is not None:
transformers_captioner(args_dict, paligemma_node_encode)
# raise Exception(
# "图片批量反推任务已完成 ; Image batch reverse push task completed")
return {"ui": {"string": ["图片批量反推任务已完成 ; Image batch reverse push task completed",]}, "result": ("", None)}
import torch
import folder_paths
from . import mz_prompt_utils
from .mz_prompt_utils import Utils
paligemma_files_map = {
"common": [
{
"url": "https://www.modelscope.cn/api/v1/models/wailovet/sd-models/repo?Revision=master&FilePath=sd3-long-captioner-merge%2Fadded_tokens.json",
"file_path": "added_tokens.json"
},
{
"url": "https://www.modelscope.cn/api/v1/models/wailovet/sd-models/repo?Revision=master&FilePath=sd3-long-captioner-merge%2Fconfig.json",
"file_path": "config.json"
},
{
"url": "https://www.modelscope.cn/api/v1/models/wailovet/sd-models/repo?Revision=master&FilePath=sd3-long-captioner-merge%2Fgeneration_config.json",
"file_path": "generation_config.json"
},
{
"url": "https://www.modelscope.cn/api/v1/models/wailovet/sd-models/repo?Revision=master&FilePath=sd3-long-captioner-merge%2Fpreprocessor_config.json",
"file_path": "preprocessor_config.json"
},
{
"url": "https://www.modelscope.cn/api/v1/models/wailovet/sd-models/repo?Revision=master&FilePath=sd3-long-captioner-merge%2Fspecial_tokens_map.json",
"file_path": "special_tokens_map.json"
},
{
"url": "https://www.modelscope.cn/api/v1/models/wailovet/sd-models/repo?Revision=master&FilePath=sd3-long-captioner-merge%2Ftokenizer.json",
"file_path": "tokenizer.json"
},
{
"url": "https://www.modelscope.cn/api/v1/models/wailovet/sd-models/repo?Revision=master&FilePath=sd3-long-captioner-merge%2Ftokenizer.model",
"file_path": "tokenizer.model"
},
{
"url": "https://www.modelscope.cn/api/v1/models/wailovet/sd-models/repo?Revision=master&FilePath=sd3-long-captioner-merge%2Ftokenizer_config.json",
"file_path": "tokenizer_config.json"
},
],
"paligemma-sd3-long-captioner": [
{
"url": "https://www.modelscope.cn/api/v1/models/wailovet/sd-models/repo?Revision=master&FilePath=sd3-long-captioner-merge%2Fmodel.safetensors",
"file_path": "model.safetensors"
},
],
"paligemma-sd3-long-captioner-v2": [
{
"url": "https://www.modelscope.cn/api/v1/models/wailovet/sd-models/repo?Revision=master&FilePath=sd3-long-captioner-v2-merge%2Fmodel.safetensors",
"file_path": "model.safetensors"
},
],
"paligemma-sdxl-long-captioner": [
{
"url": "https://www.modelscope.cn/api/v1/models/wailovet/sd-models/repo?Revision=master&FilePath=sdxl-long-captioner-merge%2Fmodel.safetensors",
"file_path": "model.safetensors"
},
],
}
llm_path = os.path.join(
folder_paths.models_dir,
"LLM",
)
os.makedirs(llm_path, exist_ok=True)
model_name = args_dict.get("model_name")
model_path = os.path.join(llm_path, model_name)
common_files = paligemma_files_map.get("common", [])
for file_info in common_files:
file_path = os.path.join(model_path, file_info["file_path"])
if not os.path.exists(file_path):
Utils.download_file(file_info["url"], file_path)
paligemma_files = paligemma_files_map.get(model_name, [])
for file_info in paligemma_files:
file_path = os.path.join(model_path, file_info["file_path"])
if not os.path.exists(file_path):
Utils.download_file(file_info["url"], file_path)
from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Utils.cache_get(f"paligemma_model_and_opt_")
if model is None:
model = PaliGemmaForConditionalGeneration.from_pretrained(
model_path,
local_files_only=True,
torch_dtype=torch.float16,
)
model.to(device).eval()
Utils.cache_set(f"paligemma_model_and_opt_", model)
processor = PaliGemmaProcessor.from_pretrained(
model_path,
local_files_only=True,
)
tensor_image = args_dict.get("image")
pil_image = Utils.tensor2pil(tensor_image)
resolution = args_dict.get("resolution", 512)
pil_image = Utils.resize_max(
pil_image, resolution, resolution).convert("RGB")
# prefix
prompt = "caption en"
model_inputs = processor(
text=prompt, images=pil_image, return_tensors="pt").to('cuda')
input_len = model_inputs["input_ids"].shape[-1]
def modify_caption(caption: str) -> str:
"""
Removes specific prefixes from captions.
Args:
caption (str): A string containing a caption.
Returns:
str: The caption with the prefix removed if it was present.
"""
# Define the prefixes to remove
import re
prefix_substrings = [
('captured from ', ''),
('captured at ', '')
]
# Create a regex pattern to match any of the prefixes
pattern = '|'.join([re.escape(opening)
for opening, _ in prefix_substrings])
replacers = {opening: replacer for opening,
replacer in prefix_substrings}
# Function to replace matched prefix with its corresponding replacement
def replace_fn(match):
return replacers[match.group(0)]
# Apply the regex to the caption
return re.sub(pattern, replace_fn, caption, count=1, flags=re.IGNORECASE)
with torch.inference_mode():
generation = model.generate(
**model_inputs, max_new_tokens=256, do_sample=False)
generation = generation[0][input_len:]
decoded = processor.decode(generation, skip_special_tokens=True)
modified_caption = modify_caption(decoded)
# print(modified_caption)
response = modified_caption
keep_device = args_dict.get("keep_device", False)
if not keep_device:
model.cpu()
del model
torch.cuda.empty_cache()
Utils.cache_set(f"paligemma_model_and_opt_", None)
conditionings = None
clip = args_dict.get("clip", None)
if clip is not None:
conditionings = Utils.a1111_clip_text_encode(
clip, response, )
return {"ui": {"string": [mz_prompt_utils.Utils.to_debug_prompt(response),]}, "result": (response, conditionings)}
web/prompt_mz.js
+1
View File
@@ -92,6 +92,7 @@ const my_ui = {
case "MZ_LLamaCPPCLIPTextEncode":
case "MZ_ImageInterrogatorCLIPTextEncode":
case "MZ_Florence2CLIPTextEncode":
case "MZ_PaliGemmaCLIPTextEncode":
// Node Created
const onNodeCreated = nodeType.prototype.onNodeCreated;
nodeType.prototype.onNodeCreated = function () {