3.1.9 2024.10.28
1. （聚合推文）MJ反推、SD反推 添加剪映分镜 2. （聚合推文）完善SD反推分类（界面同MJ反推，些许不一致） 3. （聚合推文）完成一键合成视频（单个和批量） 4. 修改聚合推文进入界面小说批次任务表格样式 5. （聚合推文）完善一键重置 6. （聚合推文）完善一键删除
This commit is contained in:
@@ -0,0 +1,29 @@
|
||||
Collecting accelerate
|
||||
Downloading accelerate-1.0.1-py3-none-any.whl.metadata (19 kB)
|
||||
Requirement already satisfied: numpy<3.0.0,>=1.17 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (2.1.2)
|
||||
Requirement already satisfied: packaging>=20.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (24.1)
|
||||
Collecting psutil (from accelerate)
|
||||
Downloading psutil-6.1.0-cp37-abi3-win_amd64.whl.metadata (23 kB)
|
||||
Requirement already satisfied: pyyaml in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (6.0.2)
|
||||
Requirement already satisfied: torch>=1.10.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (2.5.0)
|
||||
Requirement already satisfied: huggingface-hub>=0.21.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (0.26.1)
|
||||
Requirement already satisfied: safetensors>=0.4.3 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (0.4.5)
|
||||
Requirement already satisfied: filelock in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from huggingface-hub>=0.21.0->accelerate) (3.16.1)
|
||||
Requirement already satisfied: fsspec>=2023.5.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from huggingface-hub>=0.21.0->accelerate) (2024.10.0)
|
||||
Requirement already satisfied: requests in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from huggingface-hub>=0.21.0->accelerate) (2.32.3)
|
||||
Requirement already satisfied: tqdm>=4.42.1 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from huggingface-hub>=0.21.0->accelerate) (4.66.5)
|
||||
Requirement already satisfied: typing-extensions>=3.7.4.3 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from huggingface-hub>=0.21.0->accelerate) (4.12.2)
|
||||
Requirement already satisfied: networkx in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from torch>=1.10.0->accelerate) (3.4.2)
|
||||
Requirement already satisfied: jinja2 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from torch>=1.10.0->accelerate) (3.1.4)
|
||||
Requirement already satisfied: sympy==1.13.1 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from torch>=1.10.0->accelerate) (1.13.1)
|
||||
Requirement already satisfied: mpmath<1.4,>=1.1.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from sympy==1.13.1->torch>=1.10.0->accelerate) (1.3.0)
|
||||
Requirement already satisfied: colorama in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from tqdm>=4.42.1->huggingface-hub>=0.21.0->accelerate) (0.4.6)
|
||||
Requirement already satisfied: MarkupSafe>=2.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from jinja2->torch>=1.10.0->accelerate) (3.0.2)
|
||||
Requirement already satisfied: charset-normalizer<4,>=2 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from requests->huggingface-hub>=0.21.0->accelerate) (3.4.0)
|
||||
Requirement already satisfied: idna<4,>=2.5 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from requests->huggingface-hub>=0.21.0->accelerate) (3.10)
|
||||
Requirement already satisfied: urllib3<3,>=1.21.1 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from requests->huggingface-hub>=0.21.0->accelerate) (2.2.3)
|
||||
Requirement already satisfied: certifi>=2017.4.17 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from requests->huggingface-hub>=0.21.0->accelerate) (2024.8.30)
|
||||
Downloading accelerate-1.0.1-py3-none-any.whl (330 kB)
|
||||
Downloading psutil-6.1.0-cp37-abi3-win_amd64.whl (254 kB)
|
||||
Installing collected packages: psutil, accelerate
|
||||
Successfully installed accelerate-1.0.1 psutil-6.1.0
|
||||
@@ -0,0 +1,336 @@
|
||||
import spaces
|
||||
import gradio as gr
|
||||
from huggingface_hub import InferenceClient
|
||||
from torch import nn
|
||||
from transformers import AutoModel, AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM
|
||||
from pathlib import Path
|
||||
import torch
|
||||
import torch.amp.autocast_mode
|
||||
from PIL import Image
|
||||
import os
|
||||
import torchvision.transforms.functional as TVF
|
||||
|
||||
|
||||
CLIP_PATH = "google/siglip-so400m-patch14-384"
|
||||
CHECKPOINT_PATH = Path("cgrkzexw-599808")
|
||||
TITLE = "<h1><center>JoyCaption Alpha Two (2024-09-26a)</center></h1>"
|
||||
CAPTION_TYPE_MAP = {
|
||||
"Descriptive": [
|
||||
"Write a descriptive caption for this image in a formal tone.",
|
||||
"Write a descriptive caption for this image in a formal tone within {word_count} words.",
|
||||
"Write a {length} descriptive caption for this image in a formal tone.",
|
||||
],
|
||||
"Descriptive (Informal)": [
|
||||
"Write a descriptive caption for this image in a casual tone.",
|
||||
"Write a descriptive caption for this image in a casual tone within {word_count} words.",
|
||||
"Write a {length} descriptive caption for this image in a casual tone.",
|
||||
],
|
||||
"Training Prompt": [
|
||||
"Write a stable diffusion prompt for this image.",
|
||||
"Write a stable diffusion prompt for this image within {word_count} words.",
|
||||
"Write a {length} stable diffusion prompt for this image.",
|
||||
],
|
||||
"MidJourney": [
|
||||
"Write a MidJourney prompt for this image.",
|
||||
"Write a MidJourney prompt for this image within {word_count} words.",
|
||||
"Write a {length} MidJourney prompt for this image.",
|
||||
],
|
||||
"Booru tag list": [
|
||||
"Write a list of Booru tags for this image.",
|
||||
"Write a list of Booru tags for this image within {word_count} words.",
|
||||
"Write a {length} list of Booru tags for this image.",
|
||||
],
|
||||
"Booru-like tag list": [
|
||||
"Write a list of Booru-like tags for this image.",
|
||||
"Write a list of Booru-like tags for this image within {word_count} words.",
|
||||
"Write a {length} list of Booru-like tags for this image.",
|
||||
],
|
||||
"Art Critic": [
|
||||
"Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc.",
|
||||
"Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc. Keep it within {word_count} words.",
|
||||
"Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc. Keep it {length}.",
|
||||
],
|
||||
"Product Listing": [
|
||||
"Write a caption for this image as though it were a product listing.",
|
||||
"Write a caption for this image as though it were a product listing. Keep it under {word_count} words.",
|
||||
"Write a {length} caption for this image as though it were a product listing.",
|
||||
],
|
||||
"Social Media Post": [
|
||||
"Write a caption for this image as if it were being used for a social media post.",
|
||||
"Write a caption for this image as if it were being used for a social media post. Limit the caption to {word_count} words.",
|
||||
"Write a {length} caption for this image as if it were being used for a social media post.",
|
||||
],
|
||||
}
|
||||
|
||||
HF_TOKEN = os.environ.get("HF_TOKEN", None)
|
||||
|
||||
|
||||
class ImageAdapter(nn.Module):
|
||||
def __init__(self, input_features: int, output_features: int, ln1: bool, pos_emb: bool, num_image_tokens: int, deep_extract: bool):
|
||||
super().__init__()
|
||||
self.deep_extract = deep_extract
|
||||
|
||||
if self.deep_extract:
|
||||
input_features = input_features * 5
|
||||
|
||||
self.linear1 = nn.Linear(input_features, output_features)
|
||||
self.activation = nn.GELU()
|
||||
self.linear2 = nn.Linear(output_features, output_features)
|
||||
self.ln1 = nn.Identity() if not ln1 else nn.LayerNorm(input_features)
|
||||
self.pos_emb = None if not pos_emb else nn.Parameter(torch.zeros(num_image_tokens, input_features))
|
||||
|
||||
# Other tokens (<|image_start|>, <|image_end|>, <|eot_id|>)
|
||||
self.other_tokens = nn.Embedding(3, output_features)
|
||||
self.other_tokens.weight.data.normal_(mean=0.0, std=0.02) # Matches HF's implementation of llama3
|
||||
|
||||
def forward(self, vision_outputs: torch.Tensor):
|
||||
if self.deep_extract:
|
||||
x = torch.concat((
|
||||
vision_outputs[-2],
|
||||
vision_outputs[3],
|
||||
vision_outputs[7],
|
||||
vision_outputs[13],
|
||||
vision_outputs[20],
|
||||
), dim=-1)
|
||||
assert len(x.shape) == 3, f"Expected 3, got {len(x.shape)}" # batch, tokens, features
|
||||
assert x.shape[-1] == vision_outputs[-2].shape[-1] * 5, f"Expected {vision_outputs[-2].shape[-1] * 5}, got {x.shape[-1]}"
|
||||
else:
|
||||
x = vision_outputs[-2]
|
||||
|
||||
x = self.ln1(x)
|
||||
|
||||
if self.pos_emb is not None:
|
||||
assert x.shape[-2:] == self.pos_emb.shape, f"Expected {self.pos_emb.shape}, got {x.shape[-2:]}"
|
||||
x = x + self.pos_emb
|
||||
|
||||
x = self.linear1(x)
|
||||
x = self.activation(x)
|
||||
x = self.linear2(x)
|
||||
|
||||
# <|image_start|>, IMAGE, <|image_end|>
|
||||
other_tokens = self.other_tokens(torch.tensor([0, 1], device=self.other_tokens.weight.device).expand(x.shape[0], -1))
|
||||
assert other_tokens.shape == (x.shape[0], 2, x.shape[2]), f"Expected {(x.shape[0], 2, x.shape[2])}, got {other_tokens.shape}"
|
||||
x = torch.cat((other_tokens[:, 0:1], x, other_tokens[:, 1:2]), dim=1)
|
||||
|
||||
return x
|
||||
|
||||
def get_eot_embedding(self):
|
||||
return self.other_tokens(torch.tensor([2], device=self.other_tokens.weight.device)).squeeze(0)
|
||||
|
||||
|
||||
|
||||
# Load CLIP
|
||||
print("Loading CLIP")
|
||||
clip_processor = AutoProcessor.from_pretrained(CLIP_PATH)
|
||||
clip_model = AutoModel.from_pretrained(CLIP_PATH)
|
||||
clip_model = clip_model.vision_model
|
||||
|
||||
assert (CHECKPOINT_PATH / "clip_model.pt").exists()
|
||||
print("Loading VLM's custom vision model")
|
||||
checkpoint = torch.load(CHECKPOINT_PATH / "clip_model.pt", map_location='cpu')
|
||||
checkpoint = {k.replace("_orig_mod.module.", ""): v for k, v in checkpoint.items()}
|
||||
clip_model.load_state_dict(checkpoint)
|
||||
del checkpoint
|
||||
|
||||
clip_model.eval()
|
||||
clip_model.requires_grad_(False)
|
||||
clip_model.to("cuda")
|
||||
|
||||
|
||||
# Tokenizer
|
||||
print("Loading tokenizer")
|
||||
tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_PATH / "text_model", use_fast=True)
|
||||
assert isinstance(tokenizer, PreTrainedTokenizer) or isinstance(tokenizer, PreTrainedTokenizerFast), f"Tokenizer is of type {type(tokenizer)}"
|
||||
|
||||
# LLM
|
||||
print("Loading LLM")
|
||||
print("Loading VLM's custom text model")
|
||||
text_model = AutoModelForCausalLM.from_pretrained(CHECKPOINT_PATH / "text_model", device_map=0, torch_dtype=torch.bfloat16)
|
||||
text_model.eval()
|
||||
|
||||
# Image Adapter
|
||||
print("Loading image adapter")
|
||||
image_adapter = ImageAdapter(clip_model.config.hidden_size, text_model.config.hidden_size, False, False, 38, False)
|
||||
image_adapter.load_state_dict(torch.load(CHECKPOINT_PATH / "image_adapter.pt", map_location="cpu"))
|
||||
image_adapter.eval()
|
||||
image_adapter.to("cuda")
|
||||
|
||||
|
||||
@spaces.GPU()
|
||||
@torch.no_grad()
|
||||
def stream_chat(input_image: Image.Image, caption_type: str, caption_length: str | int, extra_options: list[str], name_input: str, custom_prompt: str) -> tuple[str, str]:
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# 'any' means no length specified
|
||||
length = None if caption_length == "any" else caption_length
|
||||
|
||||
if isinstance(length, str):
|
||||
try:
|
||||
length = int(length)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# Build prompt
|
||||
if length is None:
|
||||
map_idx = 0
|
||||
elif isinstance(length, int):
|
||||
map_idx = 1
|
||||
elif isinstance(length, str):
|
||||
map_idx = 2
|
||||
else:
|
||||
raise ValueError(f"Invalid caption length: {length}")
|
||||
|
||||
prompt_str = CAPTION_TYPE_MAP[caption_type][map_idx]
|
||||
|
||||
# Add extra options
|
||||
if len(extra_options) > 0:
|
||||
prompt_str += " " + " ".join(extra_options)
|
||||
|
||||
# Add name, length, word_count
|
||||
prompt_str = prompt_str.format(name=name_input, length=caption_length, word_count=caption_length)
|
||||
|
||||
if custom_prompt.strip() != "":
|
||||
prompt_str = custom_prompt.strip()
|
||||
|
||||
# For debugging
|
||||
print(f"Prompt: {prompt_str}")
|
||||
|
||||
# Preprocess image
|
||||
# NOTE: I found the default processor for so400M to have worse results than just using PIL directly
|
||||
#image = clip_processor(images=input_image, return_tensors='pt').pixel_values
|
||||
image = input_image.resize((384, 384), Image.LANCZOS)
|
||||
pixel_values = TVF.pil_to_tensor(image).unsqueeze(0) / 255.0
|
||||
pixel_values = TVF.normalize(pixel_values, [0.5], [0.5])
|
||||
pixel_values = pixel_values.to('cuda')
|
||||
|
||||
# Embed image
|
||||
# This results in Batch x Image Tokens x Features
|
||||
with torch.amp.autocast_mode.autocast('cuda', enabled=True):
|
||||
vision_outputs = clip_model(pixel_values=pixel_values, output_hidden_states=True)
|
||||
embedded_images = image_adapter(vision_outputs.hidden_states)
|
||||
embedded_images = embedded_images.to('cuda')
|
||||
|
||||
# Build the conversation
|
||||
convo = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful image captioner.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt_str,
|
||||
},
|
||||
]
|
||||
|
||||
# Format the conversation
|
||||
convo_string = tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = True)
|
||||
assert isinstance(convo_string, str)
|
||||
|
||||
# Tokenize the conversation
|
||||
# prompt_str is tokenized separately so we can do the calculations below
|
||||
convo_tokens = tokenizer.encode(convo_string, return_tensors="pt", add_special_tokens=False, truncation=False)
|
||||
prompt_tokens = tokenizer.encode(prompt_str, return_tensors="pt", add_special_tokens=False, truncation=False)
|
||||
assert isinstance(convo_tokens, torch.Tensor) and isinstance(prompt_tokens, torch.Tensor)
|
||||
convo_tokens = convo_tokens.squeeze(0) # Squeeze just to make the following easier
|
||||
prompt_tokens = prompt_tokens.squeeze(0)
|
||||
|
||||
# Calculate where to inject the image
|
||||
eot_id_indices = (convo_tokens == tokenizer.convert_tokens_to_ids("<|eot_id|>")).nonzero(as_tuple=True)[0].tolist()
|
||||
assert len(eot_id_indices) == 2, f"Expected 2 <|eot_id|> tokens, got {len(eot_id_indices)}"
|
||||
|
||||
preamble_len = eot_id_indices[1] - prompt_tokens.shape[0] # Number of tokens before the prompt
|
||||
|
||||
# Embed the tokens
|
||||
convo_embeds = text_model.model.embed_tokens(convo_tokens.unsqueeze(0).to('cuda'))
|
||||
|
||||
# Construct the input
|
||||
input_embeds = torch.cat([
|
||||
convo_embeds[:, :preamble_len], # Part before the prompt
|
||||
embedded_images.to(dtype=convo_embeds.dtype), # Image
|
||||
convo_embeds[:, preamble_len:], # The prompt and anything after it
|
||||
], dim=1).to('cuda')
|
||||
|
||||
input_ids = torch.cat([
|
||||
convo_tokens[:preamble_len].unsqueeze(0),
|
||||
torch.zeros((1, embedded_images.shape[1]), dtype=torch.long), # Dummy tokens for the image (TODO: Should probably use a special token here so as not to confuse any generation algorithms that might be inspecting the input)
|
||||
convo_tokens[preamble_len:].unsqueeze(0),
|
||||
], dim=1).to('cuda')
|
||||
attention_mask = torch.ones_like(input_ids)
|
||||
|
||||
# Debugging
|
||||
print(f"Input to model: {repr(tokenizer.decode(input_ids[0]))}")
|
||||
|
||||
#generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=False, suppress_tokens=None)
|
||||
#generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=True, top_k=10, temperature=0.5, suppress_tokens=None)
|
||||
generate_ids = text_model.generate(input_ids, inputs_embeds=input_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=True, suppress_tokens=None) # Uses the default which is temp=0.6, top_p=0.9
|
||||
|
||||
# Trim off the prompt
|
||||
generate_ids = generate_ids[:, input_ids.shape[1]:]
|
||||
if generate_ids[0][-1] == tokenizer.eos_token_id or generate_ids[0][-1] == tokenizer.convert_tokens_to_ids("<|eot_id|>"):
|
||||
generate_ids = generate_ids[:, :-1]
|
||||
|
||||
caption = tokenizer.batch_decode(generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)[0]
|
||||
|
||||
return prompt_str, caption.strip()
|
||||
|
||||
|
||||
with gr.Blocks() as demo:
|
||||
gr.HTML(TITLE)
|
||||
|
||||
with gr.Row():
|
||||
with gr.Column():
|
||||
input_image = gr.Image(type="pil", label="Input Image")
|
||||
|
||||
caption_type = gr.Dropdown(
|
||||
choices=["Descriptive", "Descriptive (Informal)", "Training Prompt", "MidJourney", "Booru tag list", "Booru-like tag list", "Art Critic", "Product Listing", "Social Media Post"],
|
||||
label="Caption Type",
|
||||
value="Descriptive",
|
||||
)
|
||||
|
||||
caption_length = gr.Dropdown(
|
||||
choices=["any", "very short", "short", "medium-length", "long", "very long"] +
|
||||
[str(i) for i in range(20, 261, 10)],
|
||||
label="Caption Length",
|
||||
value="long",
|
||||
)
|
||||
|
||||
extra_options = gr.CheckboxGroup(
|
||||
choices=[
|
||||
"If there is a person/character in the image you must refer to them as {name}.",
|
||||
"Do NOT include information about people/characters that cannot be changed (like ethnicity, gender, etc), but do still include changeable attributes (like hair style).",
|
||||
"Include information about lighting.",
|
||||
"Include information about camera angle.",
|
||||
"Include information about whether there is a watermark or not.",
|
||||
"Include information about whether there are JPEG artifacts or not.",
|
||||
"If it is a photo you MUST include information about what camera was likely used and details such as aperture, shutter speed, ISO, etc.",
|
||||
"Do NOT include anything sexual; keep it PG.",
|
||||
"Do NOT mention the image's resolution.",
|
||||
"You MUST include information about the subjective aesthetic quality of the image from low to very high.",
|
||||
"Include information on the image's composition style, such as leading lines, rule of thirds, or symmetry.",
|
||||
"Do NOT mention any text that is in the image.",
|
||||
"Specify the depth of field and whether the background is in focus or blurred.",
|
||||
"If applicable, mention the likely use of artificial or natural lighting sources.",
|
||||
"Do NOT use any ambiguous language.",
|
||||
"Include whether the image is sfw, suggestive, or nsfw.",
|
||||
"ONLY describe the most important elements of the image."
|
||||
],
|
||||
label="Extra Options"
|
||||
)
|
||||
|
||||
name_input = gr.Textbox(label="Person/Character Name (if applicable)")
|
||||
gr.Markdown("**Note:** Name input is only used if an Extra Option is selected that requires it.")
|
||||
|
||||
custom_prompt = gr.Textbox(label="Custom Prompt (optional, will override all other settings)")
|
||||
gr.Markdown("**Note:** Alpha Two is not a general instruction follower and will not follow prompts outside its training data well. Use this feature with caution.")
|
||||
|
||||
run_button = gr.Button("Caption")
|
||||
|
||||
with gr.Column():
|
||||
output_prompt = gr.Textbox(label="Prompt that was used")
|
||||
output_caption = gr.Textbox(label="Caption")
|
||||
|
||||
run_button.click(fn=stream_chat, inputs=[input_image, caption_type, caption_length, extra_options, name_input, custom_prompt], outputs=[output_prompt, output_caption])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
demo.launch()
|
||||
@@ -0,0 +1,423 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Use JoyCaption to caption images.
|
||||
"""
|
||||
import argparse
|
||||
import dataclasses
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
from pathlib import Path
|
||||
|
||||
import PIL.Image
|
||||
import torch
|
||||
import torch.amp
|
||||
import torchvision.transforms.functional as TVF
|
||||
from PIL import Image
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from tqdm import tqdm
|
||||
from transformers import (
|
||||
AutoTokenizer,
|
||||
LlavaForConditionalGeneration,
|
||||
PreTrainedTokenizer,
|
||||
PreTrainedTokenizerFast,
|
||||
)
|
||||
|
||||
|
||||
def none_or_type(value, desired_type):
|
||||
if value == "None":
|
||||
return None
|
||||
return desired_type(value)
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--glob", type=str, help="Glob pattern to find images")
|
||||
parser.add_argument("--filelist", type=str, help="File containing list of images")
|
||||
parser.add_argument("--prompt", type=str, help="Prompt to use")
|
||||
parser.add_argument(
|
||||
"--prompt-file", type=str, help="JSON file containing prompts to use"
|
||||
)
|
||||
parser.add_argument("--batch-size", type=int, default=1, help="Batch size")
|
||||
parser.add_argument(
|
||||
"--greedy", action="store_true", help="Use greedy decoding instead of sampling"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--temperature", type=float, default=0.6, help="Sampling temperature"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--top-p", type=lambda x: none_or_type(x, float), default=0.9, help="Top-p sampling"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--top-k", type=lambda x: none_or_type(x, int), default=None, help="Top-k sampling"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-new-tokens",
|
||||
type=int,
|
||||
default=256,
|
||||
help="Maximum length of the generated caption (in tokens)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-workers",
|
||||
type=int,
|
||||
default=4,
|
||||
help="Number of workers loading images in parallel",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="fancyfeast/llama-joycaption-alpha-two-hf-llava",
|
||||
help="Model to use",
|
||||
)
|
||||
|
||||
|
||||
PIL.Image.MAX_IMAGE_PIXELS = 933120000 # Quiets Pillow from giving warnings on really large images (WARNING: Exposes a risk of DoS from malicious images)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class Prompt:
|
||||
prompt: str
|
||||
weight: float
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def main():
|
||||
# Logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s"
|
||||
)
|
||||
|
||||
# Parse arguments
|
||||
args = parser.parse_args()
|
||||
logging.info(f"Arguments: {args}")
|
||||
|
||||
args.prompt = "Please describe the image."
|
||||
# Make sure we have a prompt or a prompt file
|
||||
prompts = parse_prompts(args.prompt, args.prompt_file)
|
||||
|
||||
args.filelist = (
|
||||
"C:\\Users\\27698\\Desktop\\node\\12\\12.txt"
|
||||
)
|
||||
# Find the images
|
||||
image_paths = find_images(args.glob, args.filelist)
|
||||
if len(image_paths) == 0:
|
||||
logging.warning("No images found")
|
||||
return
|
||||
logging.info(f"Found {len(image_paths)} images")
|
||||
|
||||
# Ignore all images that already have captions
|
||||
image_paths = [
|
||||
path for path in image_paths if not Path(path).with_suffix(".txt").exists()
|
||||
]
|
||||
|
||||
# Load JoyCaption
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model, use_fast=True)
|
||||
assert isinstance(tokenizer, PreTrainedTokenizer) or isinstance(
|
||||
tokenizer, PreTrainedTokenizerFast
|
||||
), f"Tokenizer is of type {type(tokenizer)}"
|
||||
llava_model = LlavaForConditionalGeneration.from_pretrained(
|
||||
args.model, torch_dtype="bfloat16"
|
||||
)
|
||||
assert isinstance(llava_model, LlavaForConditionalGeneration)
|
||||
|
||||
dataset = ImageDataset(
|
||||
prompts,
|
||||
image_paths,
|
||||
tokenizer,
|
||||
llava_model.config.image_token_index,
|
||||
llava_model.config.image_seq_length,
|
||||
)
|
||||
dataloader = DataLoader(
|
||||
dataset,
|
||||
collate_fn=dataset.collate_fn,
|
||||
num_workers=args.num_workers,
|
||||
shuffle=False,
|
||||
drop_last=False,
|
||||
batch_size=args.batch_size,
|
||||
)
|
||||
end_of_header_id = tokenizer.convert_tokens_to_ids("<|end_header_id|>")
|
||||
end_of_turn_id = tokenizer.convert_tokens_to_ids("<|eot_id|>")
|
||||
assert isinstance(end_of_header_id, int) and isinstance(end_of_turn_id, int)
|
||||
|
||||
pbar = tqdm(total=len(image_paths), desc="Captioning images...", dynamic_ncols=True)
|
||||
for batch in dataloader:
|
||||
vision_dtype = (
|
||||
llava_model.vision_tower.vision_model.embeddings.patch_embedding.weight.dtype
|
||||
)
|
||||
vision_device = (
|
||||
llava_model.vision_tower.vision_model.embeddings.patch_embedding.weight.device
|
||||
)
|
||||
language_device = (
|
||||
llava_model.language_model.get_input_embeddings().weight.device
|
||||
)
|
||||
|
||||
# Move to GPU
|
||||
pixel_values = batch["pixel_values"].to(vision_device, non_blocking=True)
|
||||
input_ids = batch["input_ids"].to(language_device, non_blocking=True)
|
||||
attention_mask = batch["attention_mask"].to(language_device, non_blocking=True)
|
||||
|
||||
# Normalize the image
|
||||
pixel_values = pixel_values / 255.0
|
||||
pixel_values = TVF.normalize(pixel_values, [0.5], [0.5])
|
||||
pixel_values = pixel_values.to(vision_dtype)
|
||||
|
||||
# Generate the captions
|
||||
generate_ids = llava_model.generate(
|
||||
input_ids=input_ids,
|
||||
pixel_values=pixel_values,
|
||||
attention_mask=attention_mask,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
do_sample=not args.greedy,
|
||||
suppress_tokens=None,
|
||||
use_cache=True,
|
||||
temperature=args.temperature,
|
||||
top_k=args.top_k,
|
||||
top_p=args.top_p,
|
||||
)
|
||||
|
||||
# Trim off the prompts
|
||||
assert isinstance(generate_ids, torch.Tensor)
|
||||
generate_ids = generate_ids.tolist()
|
||||
generate_ids = [
|
||||
trim_off_prompt(ids, end_of_header_id, end_of_turn_id)
|
||||
for ids in generate_ids
|
||||
]
|
||||
|
||||
# Decode the captions
|
||||
captions = tokenizer.batch_decode(
|
||||
generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False
|
||||
)
|
||||
captions = [c.strip() for c in captions]
|
||||
|
||||
for path, caption in zip(batch["paths"], captions):
|
||||
write_caption(Path(path), caption)
|
||||
|
||||
pbar.update(len(captions))
|
||||
|
||||
|
||||
def trim_off_prompt(input_ids: list[int], eoh_id: int, eot_id: int) -> list[int]:
|
||||
# Trim off the prompt
|
||||
while True:
|
||||
try:
|
||||
i = input_ids.index(eoh_id)
|
||||
except ValueError:
|
||||
break
|
||||
|
||||
input_ids = input_ids[i + 1 :]
|
||||
|
||||
# Trim off the end
|
||||
try:
|
||||
i = input_ids.index(eot_id)
|
||||
except ValueError:
|
||||
return input_ids
|
||||
|
||||
return input_ids[:i]
|
||||
|
||||
|
||||
def write_caption(image_path: Path, caption: str):
|
||||
caption_path = image_path.with_suffix(".txt")
|
||||
|
||||
try:
|
||||
f = os.open(
|
||||
caption_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL
|
||||
) # Write-only, create if not exist, fail if exists
|
||||
except FileExistsError:
|
||||
logging.warning(f"Caption file '{caption_path}' already exists")
|
||||
return
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to open caption file '{caption_path}': {e}")
|
||||
return
|
||||
|
||||
try:
|
||||
os.write(f, caption.encode("utf-8"))
|
||||
os.close(f)
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to write caption to '{caption_path}': {e}")
|
||||
return
|
||||
|
||||
|
||||
class ImageDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
prompts: list[Prompt],
|
||||
paths: list[Path],
|
||||
tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
|
||||
image_token_id: int,
|
||||
image_seq_length: int,
|
||||
):
|
||||
self.prompts = prompts
|
||||
self.paths = paths
|
||||
self.tokenizer = tokenizer
|
||||
self.image_token_id = image_token_id
|
||||
self.image_seq_length = image_seq_length
|
||||
self.pad_token_id = tokenizer.pad_token_id
|
||||
|
||||
def __len__(self):
|
||||
return len(self.paths)
|
||||
|
||||
def __getitem__(self, idx: int) -> dict:
|
||||
path = self.paths[idx]
|
||||
|
||||
# Pick a prompt
|
||||
prompt_str = random.choices(
|
||||
self.prompts, weights=[p.weight for p in self.prompts]
|
||||
)[0].prompt
|
||||
|
||||
# Preprocess image
|
||||
# NOTE: I don't use the Processor here and instead do it manually.
|
||||
# This is because in my testing a simple resize in Pillow yields higher quality results than the Processor,
|
||||
# and the Processor had some buggy behavior on some images.
|
||||
# And yes, with the so400m model, the model expects the image to be squished into a square, not padded.
|
||||
try:
|
||||
image = Image.open(path)
|
||||
if image.size != (384, 384):
|
||||
image = image.resize((384, 384), Image.LANCZOS)
|
||||
image = image.convert("RGB")
|
||||
pixel_values = TVF.pil_to_tensor(image)
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to load image '{path}': {e}")
|
||||
pixel_values = None # Will be filtered out later
|
||||
|
||||
# Build the conversation
|
||||
convo = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful image captioner.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt_str,
|
||||
},
|
||||
]
|
||||
|
||||
# Format the conversation
|
||||
convo_string = self.tokenizer.apply_chat_template(
|
||||
convo, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
assert isinstance(convo_string, str)
|
||||
|
||||
# Tokenize the conversation
|
||||
convo_tokens = self.tokenizer.encode(
|
||||
convo_string, add_special_tokens=False, truncation=False
|
||||
)
|
||||
|
||||
# Repeat the image tokens
|
||||
input_tokens = []
|
||||
for token in convo_tokens:
|
||||
if token == self.image_token_id:
|
||||
input_tokens.extend([self.image_token_id] * self.image_seq_length)
|
||||
else:
|
||||
input_tokens.append(token)
|
||||
|
||||
input_ids = torch.tensor(input_tokens, dtype=torch.long)
|
||||
attention_mask = torch.ones_like(input_ids)
|
||||
|
||||
return {
|
||||
"path": path,
|
||||
"pixel_values": pixel_values,
|
||||
"input_ids": input_ids,
|
||||
"attention_mask": attention_mask,
|
||||
}
|
||||
|
||||
def collate_fn(self, batch: list[dict]) -> dict:
|
||||
# Filter out images that failed to load
|
||||
batch = [item for item in batch if item["pixel_values"] is not None]
|
||||
|
||||
# Pad input_ids and attention_mask
|
||||
# Have to use left padding because HF's generate can't handle right padding it seems
|
||||
max_length = max(item["input_ids"].shape[0] for item in batch)
|
||||
n_pad = [max_length - item["input_ids"].shape[0] for item in batch]
|
||||
input_ids = torch.stack(
|
||||
[
|
||||
torch.nn.functional.pad(
|
||||
item["input_ids"], (n, 0), value=self.pad_token_id
|
||||
)
|
||||
for item, n in zip(batch, n_pad)
|
||||
]
|
||||
)
|
||||
attention_mask = torch.stack(
|
||||
[
|
||||
torch.nn.functional.pad(item["attention_mask"], (n, 0), value=0)
|
||||
for item, n in zip(batch, n_pad)
|
||||
]
|
||||
)
|
||||
|
||||
# Stack pixel values
|
||||
pixel_values = torch.stack([item["pixel_values"] for item in batch])
|
||||
|
||||
# Paths
|
||||
paths = [item["path"] for item in batch]
|
||||
|
||||
return {
|
||||
"paths": paths,
|
||||
"pixel_values": pixel_values,
|
||||
"input_ids": input_ids,
|
||||
"attention_mask": attention_mask,
|
||||
}
|
||||
|
||||
|
||||
def parse_prompts(prompt_str: str | None, prompt_file: str | None) -> list[Prompt]:
|
||||
if prompt_str is not None and prompt_file is not None:
|
||||
raise ValueError("Cannot specify both --prompt and --prompt-file")
|
||||
|
||||
if prompt_str is not None:
|
||||
return [Prompt(prompt=prompt_str, weight=1.0)]
|
||||
|
||||
if prompt_file is None:
|
||||
raise ValueError("Must specify either --prompt or --prompt-file")
|
||||
|
||||
data = json.loads(Path(prompt_file).read_text())
|
||||
|
||||
if not isinstance(data, list):
|
||||
raise ValueError("Expected JSON file to contain a list of prompts")
|
||||
|
||||
prompts = []
|
||||
|
||||
for item in data:
|
||||
if isinstance(item, str):
|
||||
prompts.append(Prompt(prompt=item, weight=1.0))
|
||||
elif (
|
||||
isinstance(item, dict)
|
||||
and "prompt" in item
|
||||
and "weight" in item
|
||||
and isinstance(item["prompt"], str)
|
||||
and isinstance(item["weight"], (int, float))
|
||||
):
|
||||
prompts.append(Prompt(prompt=item["prompt"], weight=item["weight"]))
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid prompt in JSON file. Should be either a string or an object with 'prompt' and 'weight' fields: {item}"
|
||||
)
|
||||
|
||||
if len(prompts) == 0:
|
||||
raise ValueError("No prompts found in JSON file")
|
||||
|
||||
if sum(p.weight for p in prompts) <= 0.0:
|
||||
raise ValueError("Prompt weights must sum to a positive number")
|
||||
|
||||
return prompts
|
||||
|
||||
|
||||
def find_images(glob: str | None, filelist: str | Path | None) -> list[Path]:
|
||||
if glob is None and filelist is None:
|
||||
raise ValueError("Must specify either --glob or --filelist")
|
||||
|
||||
paths = []
|
||||
|
||||
if glob is not None:
|
||||
paths.extend(Path(".").glob(glob))
|
||||
|
||||
if filelist is not None:
|
||||
paths.extend(
|
||||
(
|
||||
Path(line.strip())
|
||||
for line in Path(filelist).read_text().strip().splitlines()
|
||||
if line.strip() != ""
|
||||
)
|
||||
)
|
||||
|
||||
return paths
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,75 @@
|
||||
import torch
|
||||
import torch.amp
|
||||
import torchvision.transforms.functional as TVF
|
||||
from PIL import Image
|
||||
from transformers import AutoTokenizer, LlavaForConditionalGeneration
|
||||
|
||||
|
||||
IMAGE_PATH = "C:/Users/27698/Desktop/node/12/00001.png"
|
||||
PROMPT = "Write a long descriptive caption for this image in a formal tone."
|
||||
MODEL_NAME = "fancyfeast/llama-joycaption-alpha-two-hf-llava"
|
||||
|
||||
|
||||
# Load JoyCaption
|
||||
# bfloat16 is the native dtype of the LLM used in JoyCaption (Llama 3.1)
|
||||
# device_map=0 loads the model into the first GPU
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
|
||||
llava_model = LlavaForConditionalGeneration.from_pretrained(MODEL_NAME, torch_dtype="bfloat16", device_map="cuda:0")
|
||||
llava_model.eval()
|
||||
|
||||
with torch.no_grad():
|
||||
# Load and preprocess image
|
||||
# Normally you would use the Processor here, but the image module's processor
|
||||
# has some buggy behavior and a simple resize in Pillow yields higher quality results
|
||||
image = Image.open(IMAGE_PATH)
|
||||
|
||||
if image.size != (384, 384):
|
||||
image = image.resize((384, 384), Image.LANCZOS)
|
||||
|
||||
image = image.convert("RGB")
|
||||
pixel_values = TVF.pil_to_tensor(image)
|
||||
|
||||
# Normalize the image
|
||||
pixel_values = pixel_values / 255.0
|
||||
pixel_values = TVF.normalize(pixel_values, [0.5], [0.5])
|
||||
pixel_values = pixel_values.to(torch.bfloat16).unsqueeze(0)
|
||||
|
||||
# Build the conversation
|
||||
convo = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful image captioner.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": PROMPT,
|
||||
},
|
||||
]
|
||||
|
||||
# Format the conversation
|
||||
convo_string = tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
|
||||
|
||||
# Tokenize the conversation
|
||||
convo_tokens = tokenizer.encode(convo_string, add_special_tokens=False, truncation=False)
|
||||
|
||||
# Repeat the image tokens
|
||||
input_tokens = []
|
||||
for token in convo_tokens:
|
||||
if token == llava_model.config.image_token_index:
|
||||
input_tokens.extend([llava_model.config.image_token_index] * llava_model.config.image_seq_length)
|
||||
else:
|
||||
input_tokens.append(token)
|
||||
|
||||
input_ids = torch.tensor(input_tokens, dtype=torch.long).unsqueeze(0)
|
||||
attention_mask = torch.ones_like(input_ids)
|
||||
|
||||
# Generate the caption
|
||||
generate_ids = llava_model.generate(input_ids=input_ids.to('cuda'), pixel_values=pixel_values.to('cuda'), attention_mask=attention_mask.to('cuda'), max_new_tokens=300, do_sample=True, suppress_tokens=None, use_cache=True)[0]
|
||||
|
||||
# Trim off the prompt
|
||||
generate_ids = generate_ids[input_ids.shape[1]:]
|
||||
|
||||
# Decode the caption
|
||||
caption = tokenizer.decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
|
||||
caption = caption.strip()
|
||||
print(caption)
|
||||
Reference in New Issue
Block a user