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