Initial upload: Gemma 4 vision encoder (569.6M, 27-layer ViT with 2D RoPE)

README.md

@@ -86,16 +86,15 @@ Unlike the audio encoder (which is identical across E2B and E4B), the vision enc
 
 ```python
 import torch
-from transformers import Gemma4VisionModel, Gemma4VisionConfig, AutoProcessor
-from safetensors.torch import load_file
+from transformers import Gemma4VisionModel, AutoProcessor
 from PIL import Image
 
-# Load vision encoder
-cfg = Gemma4VisionConfig.from_pretrained("rnagabh/gemma4-vision-encoder")
-vision_model = Gemma4VisionModel(cfg)
-state_dict = load_file("path/to/model.safetensors")  # or download from repo
-vision_model.load_state_dict(state_dict, strict=True)
-vision_model = vision_model.to(dtype=torch.bfloat16, device="cuda")
+# Load vision encoder directly from this repo
+vision_model = Gemma4VisionModel.from_pretrained(
+    "rnagabh/gemma4-vision-encoder",
+    torch_dtype=torch.bfloat16,
+)
+vision_model.to("cuda")
 vision_model.eval()
 
 # Use the parent model's image processor for correct preprocessing
@@ -113,14 +112,14 @@ tokens_per_image = processed["num_soft_tokens_per_image"]  # for splitting batch
 with torch.no_grad():
     output = vision_model(pixel_values=pixel_values, pixel_position_ids=position_ids)
     embeddings = output.last_hidden_state  # (num_tokens, 1152)
-    
+
     # Mean-pool for a single image vector
     image_embedding = embeddings.float().mean(dim=0)  # (1152,)
 ```
 
 > **Important:** Always use `Gemma4ImageProcessor` from the parent model for preprocessing.
 > It handles resizing, patchification, position ID generation, and pixel normalization.
-> Manual patchification without this processor will produce degraded results.
+> Manual patchification without this processor will produce significantly degraded results.
 
 ## Benchmark Results (frozen 1152-dim embeddings, linear probe)
 
@@ -136,10 +135,6 @@ with torch.no_grad():
 
 Strong performance across all classes: airplane (0.98 F1), ship (0.98 F1), truck (0.97 F1), automobile (0.97 F1). Weakest class is cat (0.86 F1) — a fine-grained category that is inherently harder.
 
-> **Important:** Use the parent model's image processor (`Gemma4ImageProcessor` from `google/gemma-4-31B-it`)
-> for correct preprocessing. Manual patchification without proper resizing and position ID generation
-> will produce significantly degraded results.
-
 ## Files in This Repo
 
 | File | Description | Size |
@@ -151,9 +146,9 @@ Strong performance across all classes: airplane (0.98 F1), ship (0.98 F1), truck
 ## Limitations
 
 - **End-to-end trained for LLM decoding:** The encoder was trained to produce features for Gemma 4's text decoder. The 1152-dim output is the pure vision representation; the `embed_vision` projection maps to the 31B's text hidden space (5376-dim).
-- **Requires pre-patchified input:** Unlike standard ViT models that accept raw `(B, C, H, W)` images, this model expects pre-patchified `(B, num_patches, 768)` tensors with explicit position IDs.
-- **Variable aspect ratio support:** The 2D position embeddings enable non-square images, but you must provide correct `pixel_position_ids` for each patch.
-- **No built-in image preprocessing:** You need to handle resizing, normalization (the model does `2*(x-0.5)` internally), and patchification yourself, or use the parent model's processor.
+- **Requires parent model's image processor:** Use `Gemma4ImageProcessor` from `google/gemma-4-31B-it` for preprocessing. The model expects pre-patchified `(B, num_patches, 768)` tensors with explicit 2D position IDs — the processor handles this automatically.
+- **Variable aspect ratio support:** The 2D position embeddings enable non-square images. The processor generates correct position IDs for any aspect ratio.
+- **Output shape note:** The pooler strips padding and collapses the batch dimension, returning `(num_valid_tokens, 1152)`. For batched inference, use `num_soft_tokens_per_image` from the processor to split the output back into per-image embeddings.
 
 ## Extraction Details