Create caption_generator.py

caption_generator.py

@@ -0,0 +1,216 @@
+'''
+python caption_generator.py /path/to/input/image.jpg /path/to/output/directory --caption_type "Descriptive" --caption_length "long" --extra_options 0 2 5 --name_input "John"
+'''
+
+import argparse
+from pathlib import Path
+import torch
+from torch import nn
+from transformers import AutoModel, AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM
+from PIL import Image
+import torchvision.transforms.functional as TVF
+
+# Constants
+CLIP_PATH = "google/siglip-so400m-patch14-384"
+CHECKPOINT_PATH = Path("cgrkzexw-599808")
+
+# Extra options with IDs for easy selection
+EXTRA_OPTIONS = [
+    "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."
+]
+
+# Image Adapter
+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))
+        self.other_tokens = nn.Embedding(3, output_features)
+        self.other_tokens.weight.data.normal_(mean=0.0, std=0.02)
+
+    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)
+        else:
+            x = vision_outputs[-2]
+        x = self.ln1(x)
+        if self.pos_emb is not None:
+            x = x + self.pos_emb
+        x = self.linear1(x)
+        x = self.activation(x)
+        x = self.linear2(x)
+        other_tokens = self.other_tokens(torch.tensor([0, 1], device=self.other_tokens.weight.device).expand(x.shape[0], -1))
+        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 models
+def load_models():
+    print("Loading CLIP")
+    clip_processor = AutoProcessor.from_pretrained(CLIP_PATH)
+    clip_model = AutoModel.from_pretrained(CLIP_PATH)
+    clip_model = clip_model.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)
+    clip_model.eval()
+    clip_model.requires_grad_(False)
+    clip_model.to("cuda")
+
+    print("Loading tokenizer")
+    tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_PATH / "text_model", use_fast=True)
+
+    print("Loading LLM")
+    text_model = AutoModelForCausalLM.from_pretrained(CHECKPOINT_PATH / "text_model", device_map=0, torch_dtype=torch.bfloat16)
+    text_model.eval()
+
+    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")
+
+    return clip_processor, clip_model, tokenizer, text_model, image_adapter
+
+# Generate caption
+@torch.no_grad()
+def generate_caption(input_image: Image.Image, caption_type: str, caption_length: str | int, extra_options: list[str], name_input: str, custom_prompt: str, clip_processor, clip_model, tokenizer, text_model, image_adapter):
+    torch.cuda.empty_cache()
+
+    # Build prompt
+    length = None if caption_length == "any" else caption_length
+    if isinstance(length, str):
+        try:
+            length = int(length)
+        except ValueError:
+            pass
+    map_idx = 0 if length is None else 1 if isinstance(length, int) else 2
+    prompt_str = CAPTION_TYPE_MAP[caption_type][map_idx]
+
+    if len(extra_options) > 0:
+        prompt_str += " " + " ".join(extra_options)
+    prompt_str = prompt_str.format(name=name_input, length=caption_length, word_count=caption_length)
+
+    if custom_prompt.strip() != "":
+        prompt_str = custom_prompt.strip()
+
+    # Preprocess image
+    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
+    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 conversation
+    convo = [
+        {"role": "system", "content": "You are a helpful image captioner."},
+        {"role": "user", "content": prompt_str},
+    ]
+    convo_string = tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
+    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)
+    convo_tokens = convo_tokens.squeeze(0)
+    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()
+    preamble_len = eot_id_indices[1] - prompt_tokens.shape[0]
+
+    # 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],
+        embedded_images.to(dtype=convo_embeds.dtype),
+        convo_embeds[:, preamble_len:],
+    ], dim=1).to('cuda')
+
+    input_ids = torch.cat([
+        convo_tokens[:preamble_len].unsqueeze(0),
+        torch.zeros((1, embedded_images.shape[1]), dtype=torch.long),
+        convo_tokens[preamble_len:].unsqueeze(0),
+    ], dim=1).to('cuda')
+    attention_mask = torch.ones_like(input_ids)
+
+    # Generate caption
+    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)
+    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()
+
+# Main function
+def main():
+    parser = argparse.ArgumentParser(description="Generate a caption for an image.")
+    parser.add_argument("input_image", type=str, help="Path to the input image")
+    parser.add_argument("output_path", type=str, help="Path to save the output caption and image")
+    parser.add_argument("--caption_type", type=str, default="Descriptive", choices=CAPTION_TYPE_MAP.keys(), help="Type of caption to generate")
+    parser.add_argument("--caption_length", type=str, default="long", help="Length of the caption")
+    parser.add_argument("--extra_options", nargs="*", type=int, default=[], help="Extra options for caption generation (provide IDs separated by spaces)")
+    parser.add_argument("--name_input", type=str, default="", help="Name of the person/character in the image (if applicable)")
+    parser.add_argument("--custom_prompt", type=str, default="", help="Custom prompt to override default settings")
+    args = parser.parse_args()
+
+    # Map extra option IDs to their corresponding strings
+    selected_extra_options = [EXTRA_OPTIONS[i] for i in args.extra_options]
+
+    # Load models
+    clip_processor, clip_model, tokenizer, text_model, image_adapter = load_models()
+
+    # Open the input image
+    input_image = Image.open(args.input_image)
+
+    # Generate caption
+    prompt_str, caption = generate_caption(input_image, args.caption_type, args.caption_length, selected_extra_options, args.name_input, args.custom_prompt, clip_processor, clip_model, tokenizer, text_model, image_adapter)
+
+    # Save caption and image
+    output_path = Path(args.output_path)
+    output_path.mkdir(parents=True, exist_ok=True)
+    image_name = Path(args.input_image).name.replace(" ", "_")
+    output_image_path = output_path / image_name
+    input_image.save(output_image_path)
+
+    txt_file_path = output_path / f"{output_image_path.stem}.txt"
+    with open(txt_file_path, "w") as f:
+        f.write(f"Prompt: {prompt_str}\n\nCaption: {caption}")
+
+    print(f"Caption saved to {txt_file_path}")
+
+if __name__ == "__main__":
+    # Print extra options with IDs for reference
+    print("Extra Options:")
+    for i, option in enumerate(EXTRA_OPTIONS):
+        print(f"{i}: {option}")
+    main()