from transformers import AutoProcessor, AutoModelForCausalLM, BitsAndBytesConfig import torch from PIL import Image import requests import traceback import os from huggingface_hub import login login(token=os.getenv("HF_TOKEN")) class Image2Text: def __init__(self): # Load the GIT coco model preprocessor_git_large_coco = AutoProcessor.from_pretrained("microsoft/git-large-coco") model_git_large_coco = AutoModelForCausalLM.from_pretrained("microsoft/git-large-coco") self.device = "cuda" if torch.cuda.is_available() else "cpu" self.preprocessor = preprocessor_git_large_coco self.model = model_git_large_coco self.model.to(self.device) def image_description( self, image_url, max_length=50, temperature=0.1, use_sample_image=False, ): """ Generate captions for the given image. ----- Parameters image_url: Image URL The image to generate captions for. max_length: int The max length of the generated descriptions. ----- Returns str The generated image description. """ caption_git_large_coco = "" if use_sample_image: image_url = "http://images.cocodataset.org/val2017/000000039769.jpg" image = Image.open(requests.get(image_url, stream=True).raw) # Generate captions for the image using the GIT coco model try: caption_git_large_coco = self._generate_description(image, max_length, False).strip() return caption_git_large_coco except Exception as e: print(e) traceback.print_exc() def _generate_description( self, image, max_length=50, use_float_16=False, ): """ Generate captions for the given image. ----- Parameters image: PIL.Image The image to generate captions for. max_length: int The max length of the generated descriptions. use_float_16: bool Whether to use float16 precision. This can speed up inference, but may lead to worse results. ----- Returns str The generated caption. """ # inputs = preprocessor(image, return_tensors="pt").to(device) pixel_values = self.preprocessor(images=image, return_tensors="pt").pixel_values.to(self.device) generated_ids = self.model.generate( pixel_values=pixel_values, max_length=max_length, ) generated_caption = self.preprocessor.batch_decode(generated_ids, skip_special_tokens=True)[0] return generated_caption import json from pprint import pprint import bitsandbytes as bnb import pandas as pd import torch import torch.nn as nn import transformers from datasets import load_dataset from huggingface_hub import notebook_login from peft import ( LoraConfig , PeftConfig , PeftModel , get_peft_model , prepare_model_for_kbit_training, ) from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, AutoConfig from peft import LoraConfig, get_peft_model os.environ["CUDA_VISIBLE_DEVICES"] = "0" class Social_Media_Captioner: def __init__(self, use_finetuned: bool=True, temp=0.1): self.use_finetuned = use_finetuned self.MODEL_NAME = "vilsonrodrigues/falcon-7b-instruct-sharded" self.peft_model_name = "ayush-vatsal/caption_qlora_finetune" self.model_loaded = False self.device = "cuda:0" self._load_model() self.generation_config = self.model.generation_config self.generation_config.max_new_tokens = 50 self.generation_config.temperature = temp self.generation_config.top_p = 0.7 self.generation_config.num_return_sequences = 1 self.generation_config.pad_token_id = self.tokenizer.eos_token_id self.generation_config.eos_token_id = self.tokenizer.eos_token_id self.cache: list[dict] = [] # [{"image_decription": "A man", "caption": ["A man"]}] def _load_model(self): self.bnb_config = BitsAndBytesConfig( load_in_4bit = True, llm_int8_enable_fp32_cpu_offload=True, bnb_4bit_use_double_quant = True, bnb_4bit_quant_type= "nf4", bnb_4bit_compute_dtype=torch.bfloat16, load_in_8bit_fp32_cpu_offload=True ) self.model = AutoModelForCausalLM.from_pretrained( self.MODEL_NAME, device_map = "auto", trust_remote_code = True, quantization_config = self.bnb_config ) # Defining the tokenizers self.tokenizer = AutoTokenizer.from_pretrained(self.MODEL_NAME) self.tokenizer.pad_token = self.tokenizer.eos_token # if self.use_finetuned: # # LORA Config Model # self.lora_config = LoraConfig( # r=16, # lora_alpha=32, # target_modules=["query_key_value"], # lora_dropout=0.05, # bias="none", # task_type="CAUSAL_LM" # ) # self.model = get_peft_model(self.model, self.lora_config) # # Fitting the adapters # self.peft_config = PeftConfig.from_pretrained(self.peft_model_name) # self.model = AutoModelForCausalLM.from_pretrained( # self.peft_config.base_model_name_or_path, # return_dict = True, # quantization_config = self.bnb_config, # device_map= "auto", # trust_remote_code = True # ) # self.model = PeftModel.from_pretrained(self.model, self.peft_model_name) # # Defining the tokenizers # self.tokenizer = AutoTokenizer.from_pretrained(self.peft_config.base_model_name_or_path) # self.tokenizer.pad_token = self.tokenizer.eos_token self.model_loaded = True print("Model Loaded successfully") def inference(self, input_text: str, use_cached=True, cache_generation=True) -> str | None: if not self.model_loaded: raise Exception("Model not loaded") try: prompt = Social_Media_Captioner._prompt(input_text) if use_cached: for item in self.cache: if item['image_description'] == input_text: return item['caption'] encoding = self.tokenizer(prompt, return_tensors = "pt").to(self.device) with torch.inference_mode(): outputs = self.model.generate( input_ids = encoding.input_ids, attention_mask = encoding.attention_mask, generation_config = self.generation_config ) generated_caption = (self.tokenizer.decode(outputs[0], skip_special_tokens=True).split('Caption: "')[-1]).split('"')[0] if cache_generation: for item in self.cache: if item['image_description'] == input_text: item['caption'].append(generated_caption) break else: self.cache.append({ 'image_description': input_text, 'caption': [generated_caption] }) return generated_caption except Exception as e: print(e) return None def _prompt(input_text="A man walking alone in the road"): if input_text is None: raise Exception("Enter a valid input text to generate a valid prompt") return f""" Convert the given image description to a appropriate metaphoric caption Description: {input_text} Caption: """.strip() @staticmethod def get_trainable_parameters(model): trainable_params = 0 all_param = 0 for _, param in model.named_parameters(): all_param += param.numel() if param.requires_grad: trainable_params += param.numel() return f"trainable_params: {trainable_params} || all_params: {all_param} || Percentage of trainable params: {100*trainable_params / all_param}" def __repr__(self): return f""" Base Model Name: {self.MODEL_NAME} PEFT Model Name: {self.peft_model_name} Using PEFT Finetuned Model: {self.use_finetuned} Model: {self.model} ------------------------------------------------------------ {Social_Media_Captioner.get_trainable_parameters(self.model)} """ class Captions: def __init__(self, use_finetuned_LLM: bool=True, temp_LLM=0.1): self.image_to_text = Image2Text() self.LLM = Social_Media_Captioner(use_finetuned_LLM, temp_LLM) def generate_captions( self, image, image_url=None, max_length_GIT=50, temperature_GIT=0.1, use_sample_image_GIT=False, use_cached_LLM=True, cache_generation_LLM=True ): if image_url: image_description = self.image_to_text.image_description(image_url, max_length=max_length_GIT, temperature=temperature_GIT, use_sample_image=use_sample_image_GIT) else: image_description = self.image_to_text._generate_description(image, max_length=max_length_GIT) captions = self.LLM.inference(image_description, use_cached=use_cached_LLM, cache_generation=cache_generation_LLM) return captions caption_generator = Captions() import gradio as gr def setup(image): return caption_generator.generate_captions(image = image) iface = gr.Interface( fn=setup, inputs=gr.inputs.Image(type="pil", label="Upload Image"), outputs=gr.outputs.Textbox(label="Caption") ) iface.launch()