app.py

import subprocess
import sys

def install(package):
    subprocess.check_call([sys.executable, "-m", "pip", "install", package])

install("evaluate")
install("jiwer")
install("huggingface_hub")
install("gradio")
install("bitsandbytes")
install("git+https://github.com/huggingface/transformers.git")
install("git+https://github.com/huggingface/peft.git")
install("git+https://github.com/huggingface/accelerate.git")
install("einops")
install("safetensors")
install("torch")
install("xformers")
install("datasets")

from transformers import AutoProcessor, AutoModelForCausalLM, BitsAndBytesConfig
import torch
from PIL import Image
import requests
import traceback

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
import os
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):
        try:
            self.bnb_config = BitsAndBytesConfig(
                load_in_4bit = True,
                bnb_4bit_use_double_quant = True,
                bnb_4bit_quant_type= "nf4",
                bnb_4bit_compute_dtype=torch.bfloat16,
                )
            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")

        except Exception as e:
            print(e)
            self.model_loaded = False


    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 social media worthy 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):
    # Assuming `caption_generator.generate_captions` is your function to generate captions.
    # This is just a placeholder for your actual caption generation logic.
    return caption_generator.generate_captions(image = image)

iface = gr.Interface(
    fn=setup,
    inputs=gr.Image(type="pil", label="Upload Image"),  # Updated usage here
    outputs="text"  # Simplified usage here
)

iface.launch()