161 lines
4.5 KiB
Python
161 lines
4.5 KiB
Python
from folder_paths import (
|
|
get_input_directory,
|
|
get_annotated_filepath,
|
|
exists_annotated_filepath,
|
|
models_dir,
|
|
)
|
|
|
|
from os import listdir
|
|
from os.path import isfile, join, isdir, exists
|
|
from PIL import Image
|
|
|
|
from comfy.sd import CLIP
|
|
from comfy.model_management import get_torch_device, soft_empty_cache
|
|
import gc
|
|
|
|
import transformers
|
|
|
|
from preprocess import preprocess
|
|
from utility import tokenize_text
|
|
|
|
import model_utils
|
|
|
|
INT_MAX = 0xFFFFFFFFFFFFFFFF
|
|
FLOAT_MAX = 1_000_000.0
|
|
|
|
|
|
class ImageCaptionNode:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
input_dir = get_input_directory()
|
|
image_files = [f for f in listdir(input_dir) if isfile(join(input_dir, f))]
|
|
|
|
models_dir_path = join(models_dir, "image_captioners")
|
|
model_dirs = [
|
|
d for d in listdir(models_dir_path) if isdir(join(models_dir_path, d))
|
|
]
|
|
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"image": (
|
|
sorted(image_files),
|
|
{"image_upload": True},
|
|
),
|
|
"model_name": (model_dirs,),
|
|
"min_new_tokens": (
|
|
"INT",
|
|
{"default": 20, "min": 0, "max": INT_MAX, "step": 1},
|
|
),
|
|
"max_new_tokens": (
|
|
"INT",
|
|
{"default": 50, "min": 35, "max": INT_MAX, "step": 1},
|
|
),
|
|
"num_beams": (
|
|
"INT",
|
|
{"default": 20, "min": 1, "max": INT_MAX, "step": 1},
|
|
),
|
|
"penalty_alpha": (
|
|
"FLOAT",
|
|
{"default": 0.6, "min": 0.0, "max": FLOAT_MAX, "step": 0.1},
|
|
),
|
|
"top_k": ("INT", {"default": 50, "min": 0, "max": INT_MAX, "step": 1}),
|
|
"repetition_penalty": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 1.0, "max": FLOAT_MAX, "step": 0.1},
|
|
),
|
|
"preprocess_mode": (["exact_keyword", "exact_prompt", "none"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"CONDITIONING",
|
|
"STRING",
|
|
)
|
|
RETURN_NAMES = ("clip_output", "string_output")
|
|
FUNCTION = "image_caption"
|
|
CATEGORY = "image-caption"
|
|
|
|
def image_caption(
|
|
self,
|
|
clip: CLIP,
|
|
image: str,
|
|
model_name: str,
|
|
min_new_tokens: int,
|
|
max_new_tokens: int,
|
|
num_beams: int,
|
|
penalty_alpha: float,
|
|
top_k: int,
|
|
repetition_penalty: float,
|
|
preprocess_mode: str,
|
|
) -> str:
|
|
image_path = get_annotated_filepath(image)
|
|
img = Image.open(image_path).convert("RGB")
|
|
|
|
model_path = join(models_dir, "image_captioners", model_name)
|
|
|
|
dev = get_torch_device()
|
|
|
|
model = model_utils.get_model(model_path)
|
|
processor = transformers.AutoProcessor.from_pretrained(model_path)
|
|
|
|
try:
|
|
inputs = processor(
|
|
images=img,
|
|
return_tensors="pt",
|
|
padding=True,
|
|
use_fast=True,
|
|
).to(dev)
|
|
except:
|
|
inputs = processor(
|
|
images=img,
|
|
return_tensors="pt",
|
|
).to(dev)
|
|
|
|
out = model.generate(
|
|
**inputs,
|
|
num_return_sequences=1,
|
|
max_new_tokens=max_new_tokens,
|
|
min_new_tokens=min_new_tokens,
|
|
early_stopping=True,
|
|
num_beams=num_beams,
|
|
penalty_alpha=penalty_alpha,
|
|
top_k=top_k,
|
|
repetition_penalty=repetition_penalty,
|
|
remove_invalid_values=True,
|
|
renormalize_logits=True,
|
|
)
|
|
|
|
output = processor.decode(
|
|
out[0], skip_special_tokens=True, cleanup_tokenization_spaces=True
|
|
)
|
|
|
|
del model
|
|
del processor
|
|
soft_empty_cache()
|
|
gc.collect()
|
|
|
|
output = preprocess(output, preprocess_mode)
|
|
|
|
print_string = f"{' IMAGE CAPTION OUTPUT '.center(200, '#')}\n"
|
|
print_string += f"{output}\n\n"
|
|
print_string += f"{'#' * 200}\n"
|
|
|
|
print(print_string)
|
|
|
|
return (tokenize_text(clip, output), output)
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(s, image: str, model_name: str):
|
|
if not exists_annotated_filepath(image):
|
|
return f"Invalid image file: {image}"
|
|
|
|
model_path = join(models_dir, "image_captioners", model_name)
|
|
if not exists(model_path):
|
|
return f"{model_path} is not exists"
|
|
|
|
return True
|