From 06fc84c2edca8e14a707bdbe4aa5bf9fbd4db960 Mon Sep 17 00:00:00 2001 From: wailovet Date: Mon, 24 Jun 2024 00:11:38 +0800 Subject: [PATCH] 'update' --- __init__.py | 115 +++++++++++++++++++++- mz_llama_core_nodes.py | 4 +- mz_prompt_utils.py | 12 +++ mz_transformers.py | 219 +++++++++++++++++++++++++++++++++++++---- web/prompt_mz.js | 1 + 5 files changed, 329 insertions(+), 22 deletions(-) diff --git a/__init__.py b/__init__.py index 841a2ae..20d8fd7 100644 --- a/__init__.py +++ b/__init__.py @@ -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( diff --git a/mz_llama_core_nodes.py b/mz_llama_core_nodes.py index 8cf9775..4e954d4 100644 --- a/mz_llama_core_nodes.py +++ b/mz_llama_core_nodes.py @@ -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) diff --git a/mz_prompt_utils.py b/mz_prompt_utils.py index 743ce78..d3d575f 100644 --- a/mz_prompt_utils.py +++ b/mz_prompt_utils.py @@ -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") diff --git a/mz_transformers.py b/mz_transformers.py index 0c7ce68..ce2c375 100644 --- a/mz_transformers.py +++ b/mz_transformers.py @@ -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)} diff --git a/web/prompt_mz.js b/web/prompt_mz.js index 3cd960f..40c4262 100644 --- a/web/prompt_mz.js +++ b/web/prompt_mz.js @@ -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 () {