From a1566f84ddcf3435a0128dab67458b7326edcda0 Mon Sep 17 00:00:00 2001 From: wailovet Date: Wed, 26 Jun 2024 01:07:57 +0800 Subject: [PATCH 1/6] 'update' --- __init__.py | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/__init__.py b/__init__.py index 20d8fd7..2f55922 100644 --- a/__init__.py +++ b/__init__.py @@ -649,8 +649,11 @@ class MZ_Florence2Captioner: "Florence-2-large-ft", "Florence-2-large", ],), + "directory": ("STRING", {"default": "", "placeholder": "directory"}), "resolution": ("INT", {"default": 512, "min": 128, "max": 0xffffffffffffffff}), - "captioner_config": ("ImageCaptionerConfig", ), + "caption_suffix": ("STRING", {"default": ".caption"}), + "force_update": ([False, True], {"default": False}), + "prompt_fixed_beginning": ("STRING", {"default": "", }), }, "optional": { }, @@ -666,6 +669,16 @@ class MZ_Florence2Captioner: kwargs = kwargs.copy() from . import mz_transformers importlib.reload(mz_transformers) + + + kwargs["captioner_config"] = { + "directory": kwargs["directory"], + "resolution": kwargs["resolution"], + "caption_suffix": kwargs["caption_suffix"], + "force_update": kwargs["force_update"], + "prompt_fixed_beginning": kwargs["prompt_fixed_beginning"], + } + return mz_transformers.florence2_node_encode(kwargs) @@ -721,8 +734,11 @@ class MZ_PaliGemmaCaptioner: "paligemma-sd3-long-captioner-v2", "paligemma-sd3-long-captioner", ],), + "directory": ("STRING", {"default": "", "placeholder": "directory"}), "resolution": ("INT", {"default": 512, "min": 128, "max": 0xffffffffffffffff}), - "captioner_config": ("ImageCaptionerConfig", ), + "caption_suffix": ("STRING", {"default": ".caption"}), + "force_update": ([False, True], {"default": False}), + "prompt_fixed_beginning": ("STRING", {"default": "", }), }, "optional": { }, @@ -738,6 +754,13 @@ class MZ_PaliGemmaCaptioner: kwargs = kwargs.copy() from . import mz_transformers importlib.reload(mz_transformers) + kwargs["captioner_config"] = { + "directory": kwargs["directory"], + "resolution": kwargs["resolution"], + "caption_suffix": kwargs["caption_suffix"], + "force_update": kwargs["force_update"], + "prompt_fixed_beginning": kwargs["prompt_fixed_beginning"], + } return mz_transformers.paligemma_node_encode(kwargs) From 2dcee92f534ff00092c86f5c3d7a365fcefd948a Mon Sep 17 00:00:00 2001 From: wailovet Date: Sat, 29 Jun 2024 10:08:52 +0800 Subject: [PATCH 2/6] 'update' --- __init__.py | 1 - mz_transformers.py | 57 +++++++++++++++++++++++++--------------------- 2 files changed, 31 insertions(+), 27 deletions(-) diff --git a/__init__.py b/__init__.py index 2f55922..9bc2ab7 100644 --- a/__init__.py +++ b/__init__.py @@ -670,7 +670,6 @@ class MZ_Florence2Captioner: from . import mz_transformers importlib.reload(mz_transformers) - kwargs["captioner_config"] = { "directory": kwargs["directory"], "resolution": kwargs["resolution"], diff --git a/mz_transformers.py b/mz_transformers.py index ce2c375..23c0413 100644 --- a/mz_transformers.py +++ b/mz_transformers.py @@ -17,6 +17,7 @@ def transformers_captioner(args_dict, myfunc): retry_keywords = [k for k in retry_keywords if k != ""] pre_images = [] + # print("directory:", directory) for root, dirs, files in os.walk(directory): for file in files: if file.endswith(".jpg") or file.endswith(".jpeg") or file.endswith(".png"): @@ -36,40 +37,44 @@ def transformers_captioner(args_dict, myfunc): pb = mz_prompt_utils.Utils.progress_bar(len(pre_images)) for i in range(len(pre_images)): - pre_image = pre_images[i] - image_path = pre_image["image_path"] - caption_file = pre_image["caption_path"] + try: + pre_image = pre_images[i] + image_path = pre_image["image_path"] + caption_file = pre_image["caption_path"] - onec_args_dict = args_dict.copy() - del onec_args_dict["captioner_config"] + onec_args_dict = args_dict.copy() + del onec_args_dict["captioner_config"] - pil_image = Image.open(image_path) - onec_args_dict["image"] = mz_prompt_utils.Utils.pil2tensor(pil_image) + pil_image = Image.open(image_path) + onec_args_dict["image"] = mz_prompt_utils.Utils.pil2tensor( + pil_image) - if i < len(pre_images) - 1: - onec_args_dict["keep_device"] = True + if i < len(pre_images) - 1: + onec_args_dict["keep_device"] = True - thumbnail = Image.new("RGB", (pil_image.width, pil_image.height)) - thumbnail.paste(pil_image) + thumbnail = Image.new("RGB", (pil_image.width, pil_image.height)) + thumbnail.paste(pil_image) - pb.update( - i, - len(pre_images), - # 转RGB - thumbnail, - ) + pb.update( + i, + len(pre_images), + # 转RGB + thumbnail, + ) - response = myfunc(onec_args_dict) - response = response.get("result", ())[0] - response = response.strip() + response = myfunc(onec_args_dict) + response = response.get("result", ())[0] + response = response.strip() - if response != "": - with open(caption_file, "w") as f: - prompt_fixed_beginning = captioner_config.get( - "prompt_fixed_beginning", "") - f.write(prompt_fixed_beginning + response) + if response != "": + with open(caption_file, "w") as f: + prompt_fixed_beginning = captioner_config.get( + "prompt_fixed_beginning", "") + f.write(prompt_fixed_beginning + response) - result.append(response) + result.append(response) + except Exception as e: + print(f"For image {image_path}, error: {e}") return result From 647b57b55fdd2f9b3a7c1184575b2ad9c218829a Mon Sep 17 00:00:00 2001 From: wailovet Date: Thu, 4 Jul 2024 21:59:44 +0800 Subject: [PATCH 3/6] 'update' --- __init__.py | 2 + mz_transformers.py | 122 ++++++++++++++++++++++++++++++++++----------- 2 files changed, 94 insertions(+), 30 deletions(-) diff --git a/__init__.py b/__init__.py index 9bc2ab7..a62bc38 100644 --- a/__init__.py +++ b/__init__.py @@ -651,6 +651,7 @@ class MZ_Florence2Captioner: ],), "directory": ("STRING", {"default": "", "placeholder": "directory"}), "resolution": ("INT", {"default": 512, "min": 128, "max": 0xffffffffffffffff}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 0xffffffffffffffff}), "caption_suffix": ("STRING", {"default": ".caption"}), "force_update": ([False, True], {"default": False}), "prompt_fixed_beginning": ("STRING", {"default": "", }), @@ -673,6 +674,7 @@ class MZ_Florence2Captioner: kwargs["captioner_config"] = { "directory": kwargs["directory"], "resolution": kwargs["resolution"], + "batch_size": kwargs["batch_size"], "caption_suffix": kwargs["caption_suffix"], "force_update": kwargs["force_update"], "prompt_fixed_beginning": kwargs["prompt_fixed_beginning"], diff --git a/mz_transformers.py b/mz_transformers.py index 23c0413..e1de5e0 100644 --- a/mz_transformers.py +++ b/mz_transformers.py @@ -1,6 +1,7 @@ import os import shutil import subprocess +import traceback def transformers_captioner(args_dict, myfunc): @@ -11,6 +12,7 @@ def transformers_captioner(args_dict, myfunc): force_update = captioner_config.get("force_update", False) caption_suffix = captioner_config.get("caption_suffix", "") retry_keyword = captioner_config.get("retry_keyword", "") + batch_size = captioner_config.get("batch_size", 1) retry_keywords = retry_keyword.split(",") retry_keywords = [k.strip() for k in retry_keywords] @@ -36,6 +38,7 @@ def transformers_captioner(args_dict, myfunc): result = [] pb = mz_prompt_utils.Utils.progress_bar(len(pre_images)) + images_batch = [] for i in range(len(pre_images)): try: pre_image = pre_images[i] @@ -46,14 +49,30 @@ def transformers_captioner(args_dict, myfunc): del onec_args_dict["captioner_config"] pil_image = Image.open(image_path) - onec_args_dict["image"] = mz_prompt_utils.Utils.pil2tensor( - pil_image) + if len(images_batch) < batch_size: + images_batch.append({ + "image_path": image_path, + "pil_image": pil_image + }) + if i < len(pre_images) - 1: + continue if i < len(pre_images) - 1: onec_args_dict["keep_device"] = True - thumbnail = Image.new("RGB", (pil_image.width, pil_image.height)) - thumbnail.paste(pil_image) + pil_images = [] + for j in range(len(images_batch)): + pil_images.append(images_batch[j]["pil_image"]) + + # onec_args_dict["image"] = mz_prompt_utils.Utils.pil2tensor( + # pil_image) + + thumbnail = Image.new( + "RGB", (images_batch[0]["pil_image"].width * batch_size, images_batch[0]["pil_image"].height)) + + for j in range(len(images_batch)): + pil_image = images_batch[j]["pil_image"] + thumbnail.paste(pil_image, (j * pil_image.width, 0)) pb.update( i, @@ -61,20 +80,39 @@ def transformers_captioner(args_dict, myfunc): # 转RGB thumbnail, ) + onec_args_dict["images"] = pil_images + onec_args_dict["captioner_mode"] = True - response = myfunc(onec_args_dict) - response = response.get("result", ())[0] - response = response.strip() + responses = myfunc(onec_args_dict) + # print(f"responses: {responses}") + for i in range(len(images_batch)): + item = images_batch[i] + image_path = item["image_path"] + caption_file = os.path.join( + os.path.dirname(image_path), os.path.splitext(image_path)[0] + caption_suffix) + response = responses[i] + response = response.strip() - if response != "": - with open(caption_file, "w") as f: - prompt_fixed_beginning = captioner_config.get( - "prompt_fixed_beginning", "") - f.write(prompt_fixed_beginning + response) + print(f"========================{image_path}========================") + print(image_path) + print(response) + print("") + print("") + print("") - result.append(response) + + if response != "": + with open(caption_file, "w") as f: + prompt_fixed_beginning = captioner_config.get( + "prompt_fixed_beginning", "") + f.write(prompt_fixed_beginning + response) + + result.append(response) + + images_batch = [] except Exception as e: - print(f"For image {image_path}, error: {e}") + print( + f"For image {image_path}, error: {e} , stack: {traceback.format_exc()}") return result @@ -158,16 +196,28 @@ def florence2_node_encode(args_dict): local_files_only=True, trust_remote_code=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") + + captioner_mode = args_dict.get("captioner_mode", False) + if captioner_mode: + pil_images = args_dict.get("images", None) + _pil_images = [] + for pil_image in pil_images: + resolution = args_dict.get("resolution", 512) + pil_image = Utils.resize_max( + pil_image, resolution, resolution).convert("RGB") + _pil_images.append(pil_image) + pil_images = _pil_images + else: + tensor_image = args_dict.get("image", None) + pil_image = Utils.tensor2pil(tensor_image) + resolution = args_dict.get("resolution", 512) + pil_image = Utils.resize_max( + pil_image, resolution, resolution).convert("RGB") + pil_images = [pil_image] prompt = "" - - inputs = processor(text=prompt, images=pil_image, return_tensors="pt") - + prompts = [prompt for _ in pil_images] + inputs = processor(text=prompts, images=pil_images, return_tensors="pt") generated_ids = model.generate( input_ids=inputs["input_ids"].to(device), pixel_values=inputs["pixel_values"].to(device), @@ -176,15 +226,22 @@ def florence2_node_encode(args_dict): do_sample=False ) - generated_text = processor.batch_decode( - generated_ids, skip_special_tokens=False)[0] + generated_texts = processor.batch_decode( + generated_ids, skip_special_tokens=True) - parsed_answer = processor.post_process_generation( - generated_text, - task=prompt, - image_size=(pil_image.width, pil_image.height)) - # print(f"parsed_answer: {parsed_answer}") - response = parsed_answer.get(prompt) + pil_image = pil_images[0] + parsed_answers = [] + for i in range(len(generated_texts)): + generated_text = generated_texts[i] + parsed_answer = processor.post_process_generation( + generated_text, + task=prompt, + image_size=(pil_image.width, pil_image.height)) + parsed_answers.append(parsed_answer) + + response = [] + for i in range(len(parsed_answers)): + response.append(parsed_answers[i].get(prompt)) keep_device = args_dict.get("keep_device", False) if not keep_device: @@ -193,6 +250,11 @@ def florence2_node_encode(args_dict): torch.cuda.empty_cache() Utils.cache_set(f"florence_model_and_opt_", None) + if captioner_mode: + return response + else: + response = response[0] + conditionings = None clip = args_dict.get("clip", None) if clip is not None: From e29f077f29c65c914e9275980559954bc7d9b488 Mon Sep 17 00:00:00 2001 From: wailovet Date: Thu, 4 Jul 2024 23:50:41 +0800 Subject: [PATCH 4/6] 'update' --- mz_transformers.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/mz_transformers.py b/mz_transformers.py index e1de5e0..4fa904f 100644 --- a/mz_transformers.py +++ b/mz_transformers.py @@ -1,3 +1,4 @@ +import json import os import shutil import subprocess @@ -37,6 +38,8 @@ def transformers_captioner(args_dict, myfunc): result = [] + print(f"Total images: {len(pre_images)} : {json.dumps(pre_images, indent=4)}") + pb = mz_prompt_utils.Utils.progress_bar(len(pre_images)) images_batch = [] for i in range(len(pre_images)): @@ -85,15 +88,16 @@ def transformers_captioner(args_dict, myfunc): responses = myfunc(onec_args_dict) # print(f"responses: {responses}") - for i in range(len(images_batch)): - item = images_batch[i] + for j in range(len(images_batch)): + item = images_batch[j] image_path = item["image_path"] caption_file = os.path.join( os.path.dirname(image_path), os.path.splitext(image_path)[0] + caption_suffix) - response = responses[i] + response = responses[j] response = response.strip() print(f"========================{image_path}========================") + print(caption_file) print(image_path) print(response) print("") From ab92afde1340a6471b00bbde9e147411d6c5795b Mon Sep 17 00:00:00 2001 From: wailovet Date: Thu, 4 Jul 2024 23:51:47 +0800 Subject: [PATCH 5/6] 'update' --- mz_transformers.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/mz_transformers.py b/mz_transformers.py index 4fa904f..97e9e92 100644 --- a/mz_transformers.py +++ b/mz_transformers.py @@ -38,7 +38,8 @@ def transformers_captioner(args_dict, myfunc): result = [] - print(f"Total images: {len(pre_images)} : {json.dumps(pre_images, indent=4)}") + # print(f"Total images: {len(pre_images)} : {json.dumps(pre_images, indent=4)}") + print(f"Total images: {len(pre_images)}") pb = mz_prompt_utils.Utils.progress_bar(len(pre_images)) images_batch = [] @@ -96,14 +97,11 @@ def transformers_captioner(args_dict, myfunc): response = responses[j] response = response.strip() - print(f"========================{image_path}========================") - print(caption_file) + print(f"==={image_path}===") print(image_path) print(response) print("") print("") - print("") - if response != "": with open(caption_file, "w") as f: From 36a80383bbec29553edc0a08c4bd85c659b91669 Mon Sep 17 00:00:00 2001 From: wailovet Date: Thu, 4 Jul 2024 23:58:28 +0800 Subject: [PATCH 6/6] 'update' --- mz_transformers.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/mz_transformers.py b/mz_transformers.py index 97e9e92..6379dec 100644 --- a/mz_transformers.py +++ b/mz_transformers.py @@ -44,6 +44,7 @@ def transformers_captioner(args_dict, myfunc): pb = mz_prompt_utils.Utils.progress_bar(len(pre_images)) images_batch = [] for i in range(len(pre_images)): + # print(f"Processing image {i+1}/{len(pre_images)}") try: pre_image = pre_images[i] image_path = pre_image["image_path"] @@ -53,11 +54,13 @@ def transformers_captioner(args_dict, myfunc): del onec_args_dict["captioner_config"] pil_image = Image.open(image_path) + images_batch.append({ + "image_path": image_path, + "pil_image": pil_image + }) + + if len(images_batch) < batch_size: - images_batch.append({ - "image_path": image_path, - "pil_image": pil_image - }) if i < len(pre_images) - 1: continue