PaliGemma — From-Scratch Inference Engine

A from-scratch implementation of Google's PaliGemma vision-language model for multimodal inference. This project manually reconstructs the entire PaliGemma architecture — including the SigLIP vision encoder, the Gemma language decoder, and the multimodal projection layer — and loads pre-trained weights from Hugging Face safetensors checkpoints to perform image-conditioned text generation.

---

📸 Demo

Input Image — test_images/pic1.jpg

Prompt: "this building is "

Model Output:

"The building is of Indian Institute of Technology, Kharagpur. It was established in 1951."

---

🏗️ Architecture Overview

PaliGemma is a Vision-Language Model (VLM) that combines a vision encoder with a language model to generate text conditioned on both an image and a text prompt. The architecture follows a three-stage pipeline:

┌─────────────────┐     ┌───────────────────────┐     ┌──────────────────────┐
│   Input Image   │────▶│  SigLIP Vision Encoder │────▶│  Image Embeddings    │
│  (224 × 224)    │     │  (27 Transformer Layers)│    │  (256 × 1152)        │
└─────────────────┘     └───────────────────────┘     └──────────┬───────────┘
                                                                  │
                                                     ┌────────────▼───────────┐
                                                     │  Linear Projector      │
                                                     │  (1152 → 2048)         │
                                                     └────────────┬───────────┘
                                                                  │
┌─────────────────┐     ┌───────────────────────┐    ┌────────────▼───────────┐
│   Text Prompt   │────▶│  Gemma Tokenizer +    │───▶│  Merged Embeddings     │
│                 │     │  Token Embedding       │    │  (Image + Text)        │
└─────────────────┘     └───────────────────────┘    └────────────┬───────────┘
                                                                  │
                                                     ┌────────────▼───────────┐
                                                     │  Gemma Language Model   │
                                                     │  (18 Transformer Layers)│
                                                     └────────────┬───────────┘
                                                                  │
                                                     ┌────────────▼───────────┐
                                                     │  LM Head → Logits      │
                                                     │  Autoregressive Decode  │
                                                     └────────────────────────┘

Model Specifications (PaliGemma 3B — 224px)

Component	Detail
Model	paligemma-3b-pt-224
Total Parameters	~3 Billion
Vision Encoder	SigLIP-So400m/14 (27 layers, 16 heads)
Language Decoder	Gemma 2B (18 layers, 8 heads, 1 KV head)
Image Resolution	224 × 224
Patch Size	14 × 14
Image Tokens	256 (= (224/14)²)
Vocab Size	257,216
Hidden Dimension	2,048
Precision	float32

---

📂 Project Structure

PaliGemma/
├── inference.py              # Main entry point — runs autoregressive inference
├── modeling_gemma.py         # Gemma language model + PaliGemma multimodal wrapper
├── modeling_siglip.py        # SigLIP vision encoder (from scratch)
├── processing_paligemma.py   # Image preprocessing + prompt tokenization
├── utils.py                  # Loads HF safetensors weights into custom model
├── requirements.txt          # Python dependencies
├── test_images/
│   └── pic1.jpg              # Test image (IIT Kharagpur main building)
└── hf_model/                 # Pre-trained PaliGemma weights (from Hugging Face)
    ├── config.json
    ├── tokenizer.json
    ├── tokenizer.model
    ├── model-00001-of-00003.safetensors
    ├── model-00002-of-00003.safetensors
    └── model-00003-of-00003.safetensors

---

🔬 Detailed Module Breakdown

1. modeling_siglip.py — SigLIP Vision Encoder

The vision backbone is a SigLIP (Sigmoid Loss for Image-Text Pretraining) Vision Transformer. It processes the input image into a sequence of patch embeddings that the language model can attend to.

Components

Class	Purpose
SiglipVisionConfig	Configuration dataclass holding vision encoder hyperparameters (hidden size, layers, heads, patch size).
SiglipVisionEmbeddings	Splits image into 14×14 patches via a Conv2d layer, then adds learned positional embeddings.
SiglipAttention	Standard multi-head self-attention with Q/K/V projections and scaled dot-product attention.
SiglipMLP	Two-layer feed-forward network with GELU (tanh approximation) activation.
SiglipEncoderLayer	Pre-norm Transformer block: LayerNorm → Attention → Residual → LayerNorm → MLP → Residual.
SiglipEncoder	Stack of 27 SiglipEncoderLayer blocks.
SiglipVisionTransformer	Embedding → Encoder → Post-LayerNorm. Outputs [batch, 256, 1152] hidden states.
SiglipVisionModel	Top-level wrapper around SiglipVisionTransformer.

Key Details

• Patch Embedding: A Conv2d(3, 1152, kernel_size=14, stride=14) splits the 224×224 image into a 16×16 grid of patches (256 total), each projected to a 1152-dimensional vector.
• Positional Encoding: Learned absolute position embeddings for all 256 patch positions.
• Output: 256 patch embeddings of dimension 1152, passed to the multimodal projector.

---

2. modeling_gemma.py — Gemma Language Model & PaliGemma Wrapper

This is the largest module, containing the full Gemma decoder, the multimodal projection layer, and the top-level PaliGemmaForConditionalGeneration class.

Core Components

Class	Purpose
KVCache	Key-Value cache for efficient autoregressive generation (avoids recomputing past KV states).
GemmaConfig	Configuration for Gemma language model (hidden size, layers, heads, RoPE theta, etc.).
PaliGemmaConfig	Top-level config composing SiglipVisionConfig + GemmaConfig + multimodal parameters.
GemmaRMSNorm	Root Mean Square Layer Normalization (pre-norm style, no bias).
GemmaRotaryEmbedding	Computes Rotary Position Embeddings (RoPE) for position-aware attention.
GemmaMLP	Gated feed-forward network: down_proj(GELU(gate_proj(x)) * up_proj(x)).
GemmaAttention	Grouped-Query Attention (GQA) with 8 query heads and 1 KV head, plus RoPE.
GemmaDecoderLayer	Pre-norm Transformer decoder block: RMSNorm → Attention → Residual → RMSNorm → MLP → Residual.
GemmaModel	Stack of 18 GemmaDecoderLayer blocks with token embeddings and final RMSNorm.
GemmaForCausalLM	Wraps GemmaModel with a language model head for next-token prediction. Supports weight tying.
PaliGemmaMultiModalProjector	Single Linear(1152, 2048) layer projecting vision features into the language model's embedding space.
PaliGemmaForConditionalGeneration	Top-level model — orchestrates vision encoding, projection, embedding merge, and language decoding.

Grouped-Query Attention (GQA)

The Gemma model uses Grouped-Query Attention where 8 query heads share a single Key-Value head. This reduces memory and compute by a factor of 8× for KV operations while maintaining model quality. The repeat_kv function replicates KV states across query heads at attention time.

Embedding Merge (_merge_input_ids_with_image_features)

This is the critical function that fuses vision and language:

1. Creates masks to distinguish image tokens, text tokens, and padding tokens in the input sequence.
2. Places projected image features at image token positions and text embeddings at text token positions.
3. Constructs the causal attention mask and position IDs for the decoder.

---

3. processing_paligemma.py — Input Processor

Handles both image preprocessing and prompt construction.

Image Pipeline

1. Resize image to 224×224 using bicubic interpolation.
2. Rescale pixel values from [0, 255] to [0.0, 1.0] (factor = 1/255).
3. Normalize using ImageNet mean and std [0.5, 0.5, 0.5].
4. Transpose from (H, W, C) to (C, H, W) for PyTorch.

Prompt Construction

Prepends 256 <image> tokens to the text prompt:

<image><image>...<image><bos>this building is \n

This tells the model where Visual features will be injected into the sequence.

Special Tokens

• 1024 <locXXXX> tokens — for object detection tasks (bounding box coordinates).
• 128 <segXXX> tokens — for segmentation tasks.

---

4. utils.py — Model Weight Loading

Loads pre-trained weights from Hugging Face's safetensors format:

1. Loads the tokenizer using AutoTokenizer.from_pretrained().
2. Reads all .safetensors shard files, accumulating tensors into a single dictionary.
3. Parses config.json to instantiate PaliGemmaConfig.
4. Creates the PaliGemmaForConditionalGeneration model and loads weights via load_state_dict(strict=False).
5. Ties the LM head weights with the token embedding weights.

---

5. inference.py — Inference Entry Point

Implements the full autoregressive generation loop:

1. Device Selection: Automatically selects CUDA > MPS > CPU (overridable with --only_cpu).
2. Model Loading: Calls load_hf_model() to load weights and tokenizer.
3. Processor Setup: Initializes PaliGemmaProcessor with the tokenizer, image token count, and image size.
4. Generation Loop:
   • Encodes the image and prompt into model inputs.
   • Initializes an empty KVCache.
   • For each step up to max_tokens_to_generate:
     • Forward pass through PaliGemma.
     • Extract the last token's logits.
     • Either greedy decode (argmax) or sample with top-p (nucleus sampling).
     • Append generated token and break on EOS.
   • Decode tokens back to text and print.

Sampling Strategy

• Greedy (default): do_sample=False → argmax over logits.
• Top-p (nucleus): do_sample=True → softmax with temperature scaling, then cumulative probability thresholding at top_p, followed by multinomial sampling.

---

⚙️ Hardware Used

Inference was performed on an HP ProLiant DL380 Gen11 server equipped with:

Spec	Detail
GPU	2× NVIDIA H100 NVL (95 GB each)
CUDA Version	13.2

---

🚀 Getting Started

Prerequisites

• Python 3.10+
• NVIDIA GPU with CUDA support
• Access to PaliGemma weights on Hugging Face

1. Clone the Repository

git clone https://github.com/<your-username>/PaliGemma.git
cd PaliGemma

2. Install Dependencies

pip install -r requirements.txt

3. Download Model Weights

Download the PaliGemma 3B pre-trained checkpoint from Hugging Face and place it in the hf_model/ directory:

# Using the Hugging Face CLI (login required for gated model)
huggingface-cli login
huggingface-cli download google/paligemma-3b-pt-224 --local-dir ./hf_model

Your hf_model/ directory should contain:

hf_model/
├── config.json
├── tokenizer.json
├── tokenizer.model
├── tokenizer_config.json
├── special_tokens_map.json
├── preprocessor_config.json
├── generation_config.json
├── added_tokens.json
├── model-00001-of-00003.safetensors
├── model-00002-of-00003.safetensors
├── model-00003-of-00003.safetensors
└── model.safetensors.index.json

4. Run Inference

python inference.py \
    --model_path "./hf_model" \
    --prompt "this building is " \
    --image_file_path "test_images/pic1.jpg" \
    --max_tokens_to_generate 100 \
    --temperature 0.8 \
    --top_p 0.9 \
    --do_sample False \
    --only_cpu False

CLI Arguments

Argument	Type	Default	Description
--model_path	str	—	Path to the Hugging Face model directory
--prompt	str	—	Text prompt to condition generation on
--image_file_path	str	—	Path to the input image
--max_tokens_to_generate	int	100	Maximum number of tokens to generate
--temperature	float	0.8	Sampling temperature (higher = more random)
--top_p	float	0.9	Top-p (nucleus) sampling threshold
--do_sample	bool	False	Use sampling instead of greedy decoding
--only_cpu	bool	False	Force CPU-only inference

---

🔁 Inference Pipeline — Step by Step

1. Load Model         →  Read safetensors weights & config.json, instantiate PaliGemma
2. Load Tokenizer     →  AutoTokenizer from HF with special tokens (<image>, <loc>, <seg>)
3. Preprocess Image   →  Resize(224×224) → Rescale(1/255) → Normalize(μ=0.5, σ=0.5) → CHW
4. Build Prompt       →  "<image>"×256 + "<bos>" + "this building is " + "\n"
5. Tokenize           →  Convert prompt to input_ids + attention_mask
6. Vision Encode      →  SigLIP: Image → Patch Embed → 27× Transformer → [1, 256, 1152]
7. Project            →  Linear(1152 → 2048) → [1, 256, 2048]
8. Merge Embeddings   →  Replace <image> token slots with projected image features
9. Decode (loop)      →  Gemma: merged embeddings → 18× Transformer → LM Head → next token
10. Output            →  Decode token IDs to text and print

---

📚 Key Concepts

Rotary Position Embeddings (RoPE)

Instead of adding positional encodings to the input, RoPE rotates query and key vectors by position-dependent angles. This provides better extrapolation to unseen sequence lengths and makes attention inherently position-aware.

Grouped-Query Attention (GQA)

Gemma uses 8 query heads but only 1 key-value head. During attention, the single KV head is broadcast across all 8 query heads using repeat_kv(). This drastically reduces KV cache memory during autoregressive generation.

KV Cache

During autoregressive generation, previously computed key and value states are cached so each new token only requires a single forward pass through the model, rather than reprocessing the entire sequence.

Weight Tying

The language model head (lm_head) shares its weight matrix with the token embedding layer (embed_tokens). This reduces parameter count and ensures the output logit space is aligned with the input embedding space.

---

📄 License

This project uses pre-trained weights from Google's PaliGemma, which are subject to the Gemma Terms of Use. Please review and accept the terms before downloading the model weights.

---

🤝 Acknowledgements

• Google DeepMind — PaliGemma model architecture and pre-trained weights
• Hugging Face — Model hosting, safetensors format, and transformers tokenizer
• SigLIP — Sigmoid Loss for Language Image Pre-Training
• Gemma — Open-weight language models built from Gemini research