Upload 3 files

src/paligemma/model.py

@@ -0,0 +1,53 @@
+# Necessary imports
+import os
+import sys
+from dotenv import load_dotenv
+from typing import Any
+import torch
+from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor
+
+# Local imports
+from src.logger import logging
+from src.exception import CustomExceptionHandling
+
+
+# Load the Environment Variables from .env file
+load_dotenv()
+
+# Access token for using the model
+access_token = os.environ.get("ACCESS_TOKEN")
+
+
+def load_model_and_processor(model_name: str, device: str) -> Any:
+    """
+    Load the model and processor.
+
+    Args:
+        - model_name (str): The name of the model to load.
+        - device (str): The device to load the model onto.
+
+    Returns:
+        - model: The loaded model.
+        - processor: The loaded processor.
+    """
+    try:
+        # Load the model and processor
+        model = (
+            PaliGemmaForConditionalGeneration.from_pretrained(
+                model_name, torch_dtype=torch.bfloat16, token=access_token
+            )
+            .eval()
+            .to(device)
+        )
+        processor = PaliGemmaProcessor.from_pretrained(model_name, token=access_token)
+
+        # Log the successful loading of the model and processor
+        logging.info("Model and processor loaded successfully.")
+
+        # Return the model and processor
+        return model, processor
+
+    # Handle exceptions that may occur during model and processor loading
+    except Exception as e:
+        # Custom exception handling
+        raise CustomExceptionHandling(e, sys) from e

src/paligemma/response.py

@@ -0,0 +1,71 @@
+# Necessary imports
+import sys
+import PIL.Image
+import torch
+import gradio as gr
+import spaces
+
+# Local imports
+from src.config import device, model_name
+from src.paligemma.model import load_model_and_processor
+from src.logger import logging
+from src.exception import CustomExceptionHandling
+
+
+# Language dictionary
+language_dict = {
+    "English": "en",
+    "Spanish": "es",
+    "French": "fr",
+}
+
+# Model and processor
+model, processor = load_model_and_processor(model_name, device)
+
+
+@spaces.GPU
+def caption_image(image: PIL.Image.Image, max_new_tokens: int, language: str) -> str:
+    """
+    Generates a caption based on the given image using the model.
+
+    Args:
+        - image (PIL.Image.Image): The input image to be processed.
+        - max_new_tokens (int): The maximum number of new tokens to generate.
+        - language (str): The language of the generated caption.
+
+    Returns:
+        str: The generated caption text.
+    """
+    try:
+        # Check if image is None
+        if not image:
+            gr.Warning("Please provide an image.")
+
+        # Prepare the inputs
+        language = language_dict[language]
+        prompt = f"<image>caption {language}"
+        model_inputs = (
+            processor(text=prompt, images=image, return_tensors="pt")
+            .to(torch.bfloat16)
+            .to(device)
+        )
+        input_len = model_inputs["input_ids"].shape[-1]
+
+        # Generate the response
+        with torch.inference_mode():
+            generation = model.generate(
+                **model_inputs, max_new_tokens=max_new_tokens, do_sample=False
+            )
+            generation = generation[0][input_len:]
+            decoded = processor.decode(generation, skip_special_tokens=True)
+
+        # Log the successful generation of the caption
+        logging.info("Caption generated successfully.")
+
+        # Return the generated caption
+        return decoded
+
+    # Handle exceptions that may occur during caption generation
+    except Exception as e:
+        # Custom exception handling
+        raise CustomExceptionHandling(e, sys) from e