497 lines
19 KiB
Python
497 lines
19 KiB
Python
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import traceback
|
|
|
|
|
|
def transformers_captioner(args_dict, myfunc):
|
|
from . import mz_prompt_utils
|
|
import PIL.Image as Image
|
|
captioner_config = args_dict.get("captioner_config", {})
|
|
directory = captioner_config.get("directory", None)
|
|
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]
|
|
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") or file.endswith(".webp"):
|
|
image_path = os.path.join(root, file)
|
|
base_file_path = os.path.splitext(image_path)[0]
|
|
caption_file = os.path.join(
|
|
root, base_file_path + caption_suffix)
|
|
if os.path.exists(caption_file) and force_update is False:
|
|
continue
|
|
|
|
pre_images.append({
|
|
"image_path": image_path,
|
|
"caption_path": caption_file
|
|
})
|
|
|
|
result = []
|
|
|
|
# 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 = []
|
|
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"]
|
|
caption_file = pre_image["caption_path"]
|
|
|
|
onec_args_dict = args_dict.copy()
|
|
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:
|
|
if i < len(pre_images) - 1:
|
|
continue
|
|
|
|
if i < len(pre_images) - 1:
|
|
onec_args_dict["keep_device"] = True
|
|
|
|
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,
|
|
len(pre_images),
|
|
# 转RGB
|
|
thumbnail,
|
|
)
|
|
onec_args_dict["images"] = pil_images
|
|
onec_args_dict["captioner_mode"] = True
|
|
|
|
responses = myfunc(onec_args_dict)
|
|
# print(f"responses: {responses}")
|
|
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[j]
|
|
response = response.strip()
|
|
|
|
print(f"==={image_path}===")
|
|
print(image_path)
|
|
print(response)
|
|
print("")
|
|
print("")
|
|
|
|
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} , stack: {traceback.format_exc()}")
|
|
return result
|
|
|
|
|
|
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:
|
|
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)}
|
|
|
|
import torch
|
|
import folder_paths
|
|
from . import mz_prompt_utils
|
|
from .mz_prompt_utils import Utils
|
|
|
|
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"
|
|
},
|
|
|
|
# tokenizer.json
|
|
{
|
|
"file_path": "tokenizer.json",
|
|
"url": "https://www.modelscope.cn/api/v1/models/AI-ModelScope/Florence-2-large/repo?Revision=master&FilePath=tokenizer.json"
|
|
},
|
|
# tokenizer_config.json
|
|
{
|
|
"file_path": "tokenizer_config.json",
|
|
"url": "https://www.modelscope.cn/api/v1/models/AI-ModelScope/Florence-2-large/repo?Revision=master&FilePath=tokenizer_config.json"
|
|
},
|
|
# vocab.json
|
|
{
|
|
"file_path": "vocab.json",
|
|
"url": "https://www.modelscope.cn/api/v1/models/AI-ModelScope/Florence-2-large/repo?Revision=master&FilePath=vocab.json"
|
|
}
|
|
],
|
|
"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"
|
|
},
|
|
# tokenizer.json
|
|
{
|
|
"file_path": "tokenizer.json",
|
|
"url": "https://www.modelscope.cn/api/v1/models/AI-ModelScope/Florence-2-large-ft/repo?Revision=master&FilePath=tokenizer.json"
|
|
},
|
|
# tokenizer_config.json
|
|
{
|
|
"file_path": "tokenizer_config.json",
|
|
"url": "https://www.modelscope.cn/api/v1/models/AI-ModelScope/Florence-2-large-ft/repo?Revision=master&FilePath=tokenizer_config.json"
|
|
},
|
|
# vocab.json
|
|
{
|
|
"file_path": "vocab.json",
|
|
"url": "https://www.modelscope.cn/api/v1/models/AI-ModelScope/Florence-2-large-ft/repo?Revision=master&FilePath=vocab.json"
|
|
}
|
|
],
|
|
}
|
|
|
|
llm_path = os.path.join(
|
|
folder_paths.models_dir,
|
|
"LLM",
|
|
)
|
|
os.makedirs(llm_path, exist_ok=True)
|
|
|
|
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
|
|
original_env = os.environ.get("GIT_LFS_SKIP_SMUDGE")
|
|
os.environ["GIT_LFS_SKIP_SMUDGE"] = "1"
|
|
subprocess.run(
|
|
["git", "clone", "https://www.modelscope.cn/AI-ModelScope/Florence-2-large.git", model_path])
|
|
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
|
|
if not os.path.exists(file_path) or os.path.getsize(file_path) < 1024 * 1024:
|
|
Utils.download_file(file_info["url"], file_path)
|
|
|
|
# with open(os.path.join(os.path.dirname(__file__), "hook", "modeling_florence2.py"), "r") as f:
|
|
# code = f.read()
|
|
shutil.copyfile(
|
|
os.path.join(os.path.dirname(__file__),
|
|
"hook", "modeling_florence2.py"),
|
|
os.path.join(model_path, "modeling_florence2.py")
|
|
)
|
|
|
|
from transformers import AutoProcessor, AutoModelForCausalLM
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
model = Utils.cache_get(f"florence_model_and_opt_")
|
|
if model is None:
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_path,
|
|
local_files_only=True,
|
|
trust_remote_code=True,
|
|
)
|
|
model.to(device).eval()
|
|
Utils.cache_set(f"florence_model_and_opt_", model)
|
|
|
|
processor = AutoProcessor.from_pretrained(
|
|
model_path,
|
|
local_files_only=True,
|
|
trust_remote_code=True
|
|
)
|
|
|
|
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 = "<MORE_DETAILED_CAPTION>"
|
|
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),
|
|
max_new_tokens=1024,
|
|
num_beams=3,
|
|
do_sample=False
|
|
)
|
|
|
|
generated_texts = processor.batch_decode(
|
|
generated_ids, skip_special_tokens=True)
|
|
|
|
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:
|
|
model.cpu()
|
|
del model
|
|
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:
|
|
conditionings = Utils.a1111_clip_text_encode(
|
|
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,
|
|
)
|
|
|
|
|
|
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
|
|
pil_image = pil_images[0]
|
|
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]
|
|
pil_image = pil_images[0]
|
|
|
|
|
|
# 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)
|
|
|
|
if captioner_mode:
|
|
return [response]
|
|
else:
|
|
response = response[0]
|
|
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)}
|