import torch import torch.amp.autocast_mode import os import sys import logging import warnings import argparse from PIL import Image from pathlib import Path from tqdm import tqdm from torch import nn from transformers import AutoModel, AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM from typing import List, Union # Constants CLIP_PATH = "google/siglip-so400m-patch14-384" VLM_PROMPT = "A descriptive caption for this image:\n" MODEL_PATH = "unsloth/Meta-Llama-3.1-8B-bnb-4bit" CHECKPOINT_PATH = Path("wpkklhc6") IMAGE_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.bmp', '.webp') warnings.filterwarnings("ignore", category=UserWarning) logging.getLogger("transformers").setLevel(logging.ERROR) class ImageAdapter(nn.Module): def __init__(self, input_features: int, output_features: int): super().__init__() self.linear1 = nn.Linear(input_features, output_features) self.activation = nn.GELU() self.linear2 = nn.Linear(output_features, output_features) def forward(self, vision_outputs: torch.Tensor): return self.linear2(self.activation(self.linear1(vision_outputs))) def load_models(): print("Loading CLIP 📎") clip_processor = AutoProcessor.from_pretrained(CLIP_PATH) clip_model = AutoModel.from_pretrained(CLIP_PATH).vision_model.eval().requires_grad_(False).to("cuda") print("Loading tokenizer 🪙") tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_fast=False) assert isinstance(tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)), f"Tokenizer is of type {type(tokenizer)}" print("Loading LLM 🤖") text_model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map="auto", torch_dtype=torch.bfloat16).eval() print("Loading image adapter 🖼️") image_adapter = ImageAdapter(clip_model.config.hidden_size, text_model.config.hidden_size) image_adapter.load_state_dict(torch.load(CHECKPOINT_PATH / "image_adapter.pt", map_location="cpu", weights_only=True)) image_adapter.eval().to("cuda") return clip_processor, clip_model, tokenizer, text_model, image_adapter @torch.no_grad() def stream_chat(input_images: List[Image.Image], batch_size: int, pbar: tqdm, models: tuple) -> List[str]: clip_processor, clip_model, tokenizer, text_model, image_adapter = models torch.cuda.empty_cache() all_captions = [] for i in range(0, len(input_images), batch_size): batch = input_images[i:i+batch_size] try: images = clip_processor(images=batch, return_tensors='pt', padding=True).pixel_values.to('cuda') except ValueError as e: print(f"Error processing image batch: {e}") print("Skipping this batch and continuing...") continue with torch.amp.autocast_mode.autocast('cuda', enabled=True): vision_outputs = clip_model(pixel_values=images, output_hidden_states=True) image_features = vision_outputs.hidden_states[-2] embedded_images = image_adapter(image_features).to(dtype=torch.bfloat16) prompt = tokenizer.encode(VLM_PROMPT, return_tensors='pt') prompt_embeds = text_model.model.embed_tokens(prompt.to('cuda')).to(dtype=torch.bfloat16) embedded_bos = text_model.model.embed_tokens(torch.tensor([[tokenizer.bos_token_id]], device=text_model.device, dtype=torch.int64)).to(dtype=torch.bfloat16) inputs_embeds = torch.cat([ embedded_bos.expand(embedded_images.shape[0], -1, -1), embedded_images, prompt_embeds.expand(embedded_images.shape[0], -1, -1), ], dim=1).to(dtype=torch.bfloat16) input_ids = torch.cat([ torch.tensor([[tokenizer.bos_token_id]], dtype=torch.long).expand(embedded_images.shape[0], -1), torch.zeros((embedded_images.shape[0], embedded_images.shape[1]), dtype=torch.long), prompt.expand(embedded_images.shape[0], -1), ], dim=1).to('cuda') attention_mask = torch.ones_like(input_ids) generate_ids = text_model.generate( input_ids=input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=True, top_k=10, temperature=0.5, ) generate_ids = generate_ids[:, input_ids.shape[1]:] for ids in generate_ids: caption = tokenizer.decode(ids[:-1] if ids[-1] == tokenizer.eos_token_id else ids, skip_special_tokens=True, clean_up_tokenization_spaces=True) caption = caption.replace('<|end_of_text|>', '').replace('<|finetune_right_pad_id|>', '').strip() all_captions.append(caption) if pbar: pbar.update(len(batch)) return all_captions def process_directory(input_dir: Path, output_dir: Path, batch_size: int, models: tuple): output_dir.mkdir(parents=True, exist_ok=True) image_files = [f for f in input_dir.iterdir() if f.suffix.lower() in IMAGE_EXTENSIONS] images_to_process = [f for f in image_files if not (output_dir / f"{f.stem}.txt").exists()] if not images_to_process: print("No new images to process.") return with tqdm(total=len(images_to_process), desc="Processing images", unit="image") as pbar: for i in range(0, len(images_to_process), batch_size): batch_files = images_to_process[i:i+batch_size] batch_images = [Image.open(f).convert('RGB') for f in batch_files] captions = stream_chat(batch_images, batch_size, pbar, models) for file, caption in zip(batch_files, captions): with open(output_dir / f"{file.stem}.txt", 'w', encoding='utf-8') as f: f.write(caption) for img in batch_images: img.close() def parse_arguments(): parser = argparse.ArgumentParser(description="Process images and generate captions.") parser.add_argument("input", nargs='+', help="Input image file or directory (or multiple directories)") parser.add_argument("--output", help="Output directory (optional)") parser.add_argument("--bs", type=int, default=4, help="Batch size (default: 4)") return parser.parse_args() def main(): args = parse_arguments() input_paths = [Path(input_path) for input_path in args.input] batch_size = args.bs models = load_models() for input_path in input_paths: if input_path.is_file() and input_path.suffix.lower() in IMAGE_EXTENSIONS: output_path = input_path.with_suffix('.txt') print(f"Processing single image 🎞️: {input_path.name}") with tqdm(total=1, desc="Processing image", unit="image") as pbar: captions = stream_chat([Image.open(input_path).convert('RGB')], 1, pbar, models) with open(output_path, 'w', encoding='utf-8') as f: f.write(captions[0]) print(f"Output saved to {output_path}") elif input_path.is_dir(): output_path = Path(args.output) if args.output else input_path print(f"Processing directory 📁: {input_path}") print(f"Output directory 📦: {output_path}") print(f"Batch size 🗄️: {batch_size}") process_directory(input_path, output_path, batch_size, models) else: print(f"Invalid input: {input_path}") print("Skipping...") if not input_paths: print("Usage:") print("For single image: python app.py [image_file] [--bs batch_size]") print("For directory (same input/output): python app.py [directory] [--bs batch_size]") print("For directory (separate input/output): python app.py [directory] --output [output_directory] [--bs batch_size]") print("For multiple directories: python app.py [directory1] [directory2] ... [--output output_directory] [--bs batch_size]") sys.exit(1) if __name__ == "__main__": main()