diff --git a/litgpt/adapter.py b/litgpt/adapter.py index fb39feecf5..4993885d60 100644 --- a/litgpt/adapter.py +++ b/litgpt/adapter.py @@ -19,6 +19,7 @@ from litgpt.model import GPT as BaseModel from litgpt.model import Block as BaseBlock from litgpt.model import CausalSelfAttention as BaseCausalSelfAttention +from litgpt.vision import MultiModalProjector, VisionEncoder @dataclass @@ -45,6 +46,18 @@ def __init__(self, config: Config) -> None: self.mask_cache: torch.Tensor | None = None self.max_seq_length = self.config.block_size + # Optional vision encoder for multimodal models + if config.is_multimodal: + self.vision_encoder = VisionEncoder(config, pretrained_model_name=config.vision_model_name) + self.mm_projector = MultiModalProjector( + vision_dim=config.vision_feature_dim, + text_dim=config.n_embd, + projector_type=config.mm_projector_type or "linear", + ) + else: + self.vision_encoder = None + self.mm_projector = None + @classmethod def from_name(cls, name: str, **kwargs: Any) -> Self: return cls(Config.from_name(name, **kwargs)) diff --git a/litgpt/adapter_v2.py b/litgpt/adapter_v2.py index 68a6203815..9f27c2522b 100644 --- a/litgpt/adapter_v2.py +++ b/litgpt/adapter_v2.py @@ -22,6 +22,7 @@ from litgpt.model import Block as BaseBlock from litgpt.scripts.convert_hf_checkpoint import qkv_reassemble from litgpt.utils import map_old_state_dict_weights +from litgpt.vision import MultiModalProjector, VisionEncoder @dataclass @@ -80,6 +81,18 @@ def __init__(self, config: Config) -> None: self.mask_cache: torch.Tensor | None = None self.max_seq_length = self.config.block_size + # Optional vision encoder for multimodal models + if config.is_multimodal: + self.vision_encoder = VisionEncoder(config, pretrained_model_name=config.vision_model_name) + self.mm_projector = MultiModalProjector( + vision_dim=config.vision_feature_dim, + text_dim=config.n_embd, + projector_type=config.mm_projector_type or "linear", + ) + else: + self.vision_encoder = None + self.mm_projector = None + @classmethod def from_name(cls, name: str, **kwargs: Any) -> Self: return cls(Config.from_name(name, **kwargs)) diff --git a/litgpt/api.py b/litgpt/api.py index 49410286af..ed2e1615a4 100644 --- a/litgpt/api.py +++ b/litgpt/api.py @@ -33,6 +33,7 @@ load_checkpoint, save_config, ) +from litgpt.vision import ImagePreprocessor class LLM(torch.nn.Module): @@ -469,6 +470,7 @@ def generate( top_p: float = 1.0, return_as_token_ids: bool = False, stream: bool = False, + image: str | Path | None = None, ) -> str | torch.Tensor: """ Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested. @@ -499,12 +501,31 @@ def generate( stream: If True, returns a generator that yields tokens as they are generated. At the moment, this setting is slower and may use more memory than the non-streaming version. We plan to resolve this in the future. + image: Optional path to an image file for multimodal models. """ if self.model is None: raise AttributeError( "The model is not initialized yet; use the .distribute() " "or .trainer_setup() method to initialize the model." ) + + # Preprocess image if provided + pixel_values = None + if image is not None: + if not self.config.is_multimodal: + raise ValueError( + "An image was provided but the model is not multimodal. " + "Ensure the model config has vision_feature_dim set." + ) + preprocessor = ImagePreprocessor( + image_size=self.config.vision_image_size or 224, + ) + if self.fabric is not None: + device = self.fabric.device + else: + device = self.preprocessor.device + pixel_values = preprocessor(image, device=device) + input_ids = self._text_to_token_ids(prompt, sys_prompt) prompt_length = input_ids.size(0) max_returned_tokens = prompt_length + max_new_tokens @@ -559,6 +580,7 @@ def iterator(): top_p=top_p, eos_id=self.preprocessor.tokenizer.eos_id, include_prompt=False, + pixel_values=pixel_values, ) if stream: diff --git a/litgpt/chat/base.py b/litgpt/chat/base.py index e3d22bf409..c13adc0829 100644 --- a/litgpt/chat/base.py +++ b/litgpt/chat/base.py @@ -35,6 +35,7 @@ def generate( top_k: int | None = None, top_p: float = 1.0, stop_tokens: tuple[list[int], ...] = (), + pixel_values: torch.Tensor | None = None, ) -> Iterator[torch.Tensor]: """Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as possible. @@ -72,11 +73,22 @@ def generate( top_k=top_k, top_p=top_p, stop_tokens=stop_tokens, + pixel_values=pixel_values, ) def process_prompt( - prompt, model, tokenizer, prompt_style, fabric, temperature, max_new_tokens, top_k, top_p, stop_tokens + prompt, + model, + tokenizer, + prompt_style, + fabric, + temperature, + max_new_tokens, + top_k, + top_p, + stop_tokens, + pixel_values: torch.Tensor | None = None, ): prompt = prompt_style.apply(prompt=prompt) encoded_prompt = tokenizer.encode(prompt, device=fabric.device) @@ -98,6 +110,7 @@ def process_prompt( top_k=top_k, top_p=top_p, stop_tokens=stop_tokens, + pixel_values=pixel_values, ) token_generator: Iterator[str] = tokenizer.decode_stream(y, device=fabric.device) @@ -121,7 +134,30 @@ def process_prompt( fabric.print() -def interact(multiline, model, tokenizer, prompt_style, fabric, temperature, max_new_tokens, top_k, top_p, stop_tokens): +def interact( + multiline, + model, + tokenizer, + prompt_style, + fabric, + temperature, + max_new_tokens, + top_k, + top_p, + stop_tokens, + initial_image: Path | None = None, +): + pixel_values = None + if initial_image is not None: + if not model.config.is_multimodal: + fabric.print("Warning: An image was provided but the model is not multimodal.", file=sys.stderr) + else: + from litgpt.vision import ImagePreprocessor + + preprocessor = ImagePreprocessor(image_size=model.config.vision_image_size or 224) + pixel_values = preprocessor(initial_image, device=fabric.device) + fabric.print(f">> Loaded image: {initial_image}") + while True: try: if not multiline: @@ -144,7 +180,17 @@ def interact(multiline, model, tokenizer, prompt_style, fabric, temperature, max break process_prompt( - prompt, model, tokenizer, prompt_style, fabric, temperature, max_new_tokens, top_k, top_p, stop_tokens + prompt, + model, + tokenizer, + prompt_style, + fabric, + temperature, + max_new_tokens, + top_k, + top_p, + stop_tokens, + pixel_values, ) @@ -161,6 +207,7 @@ def main( compile: bool = False, multiline: bool = False, access_token: str | None = None, + image: Path | None = None, ) -> None: """Chat with a model. @@ -267,6 +314,7 @@ def main( top_k=top_k, top_p=top_p, stop_tokens=stop_tokens, + initial_image=image, ) if fabric.device.type == "cuda": diff --git a/litgpt/config.py b/litgpt/config.py index bdba7eeac7..58635b6b73 100644 --- a/litgpt/config.py +++ b/litgpt/config.py @@ -114,6 +114,17 @@ class Config: # `rope_base` is used, for 1 `rope_local_base_freq` is used. If # `len(rope_indices) > n_layer`, we only use the initial part. rope_indices: list[int] | None = None + # Vision encoder config (optional, for multimodal models) + vision_feature_dim: int | None = None # Output dim of vision encoder (e.g. 1152 for SigLIP) + vision_start_token_id: int | None = None # Token ID for placeholder + vision_patch_size: int | None = None # Patch size (e.g. 14) + vision_image_size: int | None = None # Input image size (e.g. 224) + mm_projector_type: str | None = None # "linear" or "mlp2x" + vision_model_name: str | None = None # HF model name for the vision encoder + + @property + def is_multimodal(self) -> bool: + return self.vision_feature_dim is not None def __post_init__(self): if not self.name: diff --git a/litgpt/generate/base.py b/litgpt/generate/base.py index 12f7f4394e..4dc279a423 100644 --- a/litgpt/generate/base.py +++ b/litgpt/generate/base.py @@ -79,9 +79,10 @@ def next_token( input_pos: torch.Tensor, x: torch.Tensor, input_pos_maxp1: int | None = None, + pixel_values: torch.Tensor | None = None, **sample_kwargs: dict[str, Any], ) -> torch.Tensor: - logits = model(x, input_pos, input_pos_maxp1=input_pos_maxp1) + logits = model(x, input_pos, input_pos_maxp1=input_pos_maxp1, pixel_values=pixel_values) _next = sample(logits, **sample_kwargs).to(dtype=torch.int64) return _next @@ -137,6 +138,7 @@ def generate_fn( stop_tokens: tuple[list[int], ...] = (), include_prompt: bool, include_eos: bool, + pixel_values: torch.Tensor | None = None, ) -> Iterator[torch.Tensor]: """ Generates tokens for a single prompt. @@ -182,11 +184,13 @@ def generate_fn( input_pos_maxp1 = prompt_size if all(m.__class__.__name__ != "ThunderModule" for m in model.modules()) else None for current_idx in range(max_returned_tokens - prompt_size): # Generate the token + # pixel_values is only needed for the prefill (first forward pass) token = next_token( model, input_pos, token.view(1, -1), input_pos_maxp1=input_pos_maxp1, + pixel_values=pixel_values if prefill_token else None, temperature=temperature, top_k=top_k, top_p=top_p, @@ -380,6 +384,7 @@ def generate( top_p: float = 1.0, eos_id: int | None = None, include_prompt: bool = True, + pixel_values: torch.Tensor | None = None, ) -> torch.Tensor: """ Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested. @@ -407,6 +412,7 @@ def generate( or https://huyenchip.com/2024/01/16/sampling.html#top_p eos_id: If specified, stop generating any more token once the token is triggered. include_prompt: If true (default) prepends the prompt (after applying the prompt style) to the output. + pixel_values: Optional image tensor for multimodal models. """ token_list = list( @@ -420,6 +426,7 @@ def generate( top_k=top_k, top_p=top_p, stop_tokens=(([eos_id],) if eos_id is not None else ()), + pixel_values=pixel_values, ) ) diff --git a/litgpt/lora.py b/litgpt/lora.py index 5f9788ebf0..91cf3d7cdc 100644 --- a/litgpt/lora.py +++ b/litgpt/lora.py @@ -59,6 +59,7 @@ from litgpt.model import CausalSelfAttention as BaseCausalSelfAttention from litgpt.scripts.convert_hf_checkpoint import qkv_reassemble from litgpt.utils import map_old_state_dict_weights +from litgpt.vision import MultiModalProjector, VisionEncoder class LoRALayer(nn.Module): @@ -503,6 +504,18 @@ def __init__(self, config: Config) -> None: self.mask_cache: torch.Tensor | None = None self.max_seq_length = self.config.block_size + # Optional vision encoder for multimodal models + if config.is_multimodal: + self.vision_encoder = VisionEncoder(config, pretrained_model_name=config.vision_model_name) + self.mm_projector = MultiModalProjector( + vision_dim=config.vision_feature_dim, + text_dim=config.n_embd, + projector_type=config.mm_projector_type or "linear", + ) + else: + self.vision_encoder = None + self.mm_projector = None + @classmethod def from_name(cls, name: str, **kwargs: Any) -> Self: return cls(Config.from_name(name, **kwargs)) diff --git a/litgpt/model.py b/litgpt/model.py index 541860ab5b..1810fc0a82 100644 --- a/litgpt/model.py +++ b/litgpt/model.py @@ -17,6 +17,7 @@ from litgpt.config import Config from litgpt.scripts.convert_hf_checkpoint import qkv_reassemble +from litgpt.vision import MultiModalProjector, VisionEncoder, merge_input_embeds class GPT(nn.Module): @@ -36,6 +37,18 @@ def __init__(self, config: Config) -> None: self.mask_cache: torch.Tensor | None = None self.max_seq_length = self.config.block_size + # Optional vision encoder for multimodal models + if config.is_multimodal: + self.vision_encoder = VisionEncoder(config, pretrained_model_name=config.vision_model_name) + self.mm_projector = MultiModalProjector( + vision_dim=config.vision_feature_dim, + text_dim=config.n_embd, + projector_type=config.mm_projector_type or "linear", + ) + else: + self.vision_encoder = None + self.mm_projector = None + @property def max_seq_length(self) -> int: return self._max_seq_length @@ -88,6 +101,7 @@ def forward( input_pos: torch.Tensor | None = None, input_pos_maxp1: int | None = None, lm_head_chunk_size: int = 0, + pixel_values: torch.Tensor | None = None, ) -> torch.Tensor | list[torch.Tensor]: """ If `input_pos` is provided, the KV cache uses K and V vectors for @@ -156,6 +170,12 @@ def forward( if self.config.scale_embeddings: x = x * torch.tensor(self.config.n_embd**0.5, dtype=x.dtype) + # Merge image embeddings if pixel_values are provided + if pixel_values is not None and self.vision_encoder is not None: + image_features = self.vision_encoder(pixel_values) + image_embeds = self.mm_projector(image_features) + x = merge_input_embeds(x, image_embeds, self.config.vision_start_token_id, idx) + for block_idx, block in enumerate(self.transformer.h): if self.config.rope_indices is not None: x = block( diff --git a/litgpt/scripts/convert_hf_checkpoint.py b/litgpt/scripts/convert_hf_checkpoint.py index 7fc96f7e55..a483ceb902 100644 --- a/litgpt/scripts/convert_hf_checkpoint.py +++ b/litgpt/scripts/convert_hf_checkpoint.py @@ -307,15 +307,27 @@ def copy_weights_gemma_3( else "language_model.model" ) - GEMMA3_VISION_MODEL_PREFIX = ( - "model.vision_tower" if any(k.startswith("model.vision_tower") for k in hf_weights) else "vision_tower" - ) - - GEMMA3_MM_PROJECTOR_PREFIX = ( - "model.multi_modal_projector" - if any(k.startswith("model.multi_modal_projector") for k in hf_weights) - else "multi_modal_projector" - ) + # Detect vision encoder prefix: older HF uses "vision_tower" / "model.vision_tower", + # newer Gemma3ForConditionalGeneration uses "vision_encoder" / "model.vision_encoder". + if any(k.startswith("model.vision_tower") for k in hf_weights): + GEMMA3_VISION_MODEL_PREFIX = "model.vision_tower" + elif any(k.startswith("vision_tower") for k in hf_weights): + GEMMA3_VISION_MODEL_PREFIX = "vision_tower" + elif any(k.startswith("model.vision_encoder") for k in hf_weights): + GEMMA3_VISION_MODEL_PREFIX = "model.vision_encoder" + else: + GEMMA3_VISION_MODEL_PREFIX = "vision_encoder" + + # Detect mm-projector prefix: older HF uses "multi_modal_projector" / "model.multi_modal_projector", + # newer Gemma3ForConditionalGeneration uses "mm_projector" / "model.mm_projector". + if any(k.startswith("model.multi_modal_projector") for k in hf_weights): + GEMMA3_MM_PROJECTOR_PREFIX = "model.multi_modal_projector" + elif any(k.startswith("multi_modal_projector") for k in hf_weights): + GEMMA3_MM_PROJECTOR_PREFIX = "multi_modal_projector" + elif any(k.startswith("model.mm_projector") for k in hf_weights): + GEMMA3_MM_PROJECTOR_PREFIX = "model.mm_projector" + else: + GEMMA3_MM_PROJECTOR_PREFIX = "mm_projector" weight_map = { "model.embed_tokens.weight": "transformer.wte.weight", @@ -338,10 +350,15 @@ def copy_weights_gemma_3( if progress_per_file is not None: progress_per_file = progress_per_file / max(1, len(hf_weights) + len(qkv_weights)) - # gemma3 4b+ are multimodel models, but we are only loading the text weights + # gemma3 4b+ are multimodal models. We only load the vision tower / mm-projector weights + # when the litgpt config explicitly declares a matching vision architecture (config.is_multimodal); + # litgpt's generic vision encoder/projector shapes don't match Gemma3's SigLIP tower and + # Gemma3MultiModalProjector 1:1, so loading them unconditionally would corrupt the state_dict. is_multimodal = any(k.startswith(GEMMA3_LANGUAGE_MODEL_PREFIX) for k in hf_weights) + load_vision_weights = is_multimodal and config is not None and config.is_multimodal if is_multimodal: - warnings.warn("For Gemma3 models only the text component is supported.") + if not load_vision_weights: + warnings.warn("For Gemma3 models only the text component is supported.") new_weight_map = dict() prefix = "model" for k, v in weight_map.items(): @@ -350,8 +367,27 @@ def copy_weights_gemma_3( new_weight_map[k] = v weight_map = new_weight_map for from_name, param in hf_weights.items(): - if from_name.startswith(GEMMA3_VISION_MODEL_PREFIX) or from_name.startswith(GEMMA3_MM_PROJECTOR_PREFIX): + if from_name.startswith(GEMMA3_VISION_MODEL_PREFIX): + if load_vision_weights: + to_name = from_name.replace(GEMMA3_VISION_MODEL_PREFIX, "vision_encoder._encoder") + param = load_param(param, from_name, dtype, verbose=debug_mode) + if saver is not None: + param = saver.store_early(param) + state_dict[to_name] = param + if progress_per_file is not None: + pbar.update(progress_per_file) + continue + if from_name.startswith(GEMMA3_MM_PROJECTOR_PREFIX): + if load_vision_weights: + to_name = from_name.replace(GEMMA3_MM_PROJECTOR_PREFIX, "mm_projector.proj") + param = load_param(param, from_name, dtype, verbose=debug_mode) + if saver is not None: + param = saver.store_early(param) + state_dict[to_name] = param + if progress_per_file is not None: + pbar.update(progress_per_file) continue + name_template, *ids = layer_template(from_name, num_matches=2) to_name = weight_map.get(name_template) param = load_param(param, from_name, dtype, verbose=debug_mode) diff --git a/litgpt/vision.py b/litgpt/vision.py new file mode 100644 index 0000000000..14c1ab6c37 --- /dev/null +++ b/litgpt/vision.py @@ -0,0 +1,253 @@ +# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file. +# +# Vision encoder and multimodal projection utilities for VLMs. + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn + +from litgpt.config import Config + + +class VisionEncoder(nn.Module): + """Wraps a pretrained vision backbone (CLIP, SigLIP, etc.) and extracts patch features. + + The encoder is kept **frozen** by default – only the projector is trainable. + + Args: + config: The LitGPT model config (must have ``vision_feature_dim`` set). + pretrained_model_name: Optional HuggingFace model name. If provided, + loads weights from HuggingFace ``transformers``. + """ + + def __init__(self, config: Config, pretrained_model_name: str | None = None) -> None: + super().__init__() + if config.vision_feature_dim is None: + raise ValueError("VisionEncoder requires config.vision_feature_dim to be set.") + + self.config = config + self.vision_feature_dim = config.vision_feature_dim + self._encoder: nn.Module | None = None + + if pretrained_model_name is not None: + self._load_hf_encoder(pretrained_model_name) + else: + # Placeholder linear for testing / when loading weights separately + image_size = config.vision_image_size or 224 + patch_size = config.vision_patch_size or 14 + num_patches = (image_size // patch_size) ** 2 + # Simple conv-based patch embedding as fallback + self.patch_embed = nn.Conv2d( + 3, + self.vision_feature_dim, + kernel_size=patch_size, + stride=patch_size, + bias=False, + ) + self._num_patches = num_patches + + def _load_hf_encoder(self, model_name: str) -> None: + """Load a vision encoder from HuggingFace transformers.""" + try: + from transformers import AutoModel + except ImportError: + raise ImportError( + "Loading a pretrained vision encoder requires `transformers`. Install it with: pip install transformers" + ) + self._encoder = AutoModel.from_pretrained(model_name) + # Freeze the vision encoder + for param in self._encoder.parameters(): + param.requires_grad = False + + @property + def num_patches(self) -> int: + """Number of image patch tokens produced per image.""" + if self._encoder is not None: + # Try to get from config + image_size = self.config.vision_image_size or 224 + patch_size = self.config.vision_patch_size or 14 + return (image_size // patch_size) ** 2 + return self._num_patches + + @torch.no_grad() + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + """ + Args: + pixel_values: ``(B, C, H, W)`` image tensor, pre-normalized. + + Returns: + Image features of shape ``(B, num_patches, vision_feature_dim)``. + """ + if self._encoder is not None: + outputs = self._encoder(pixel_values=pixel_values) + # Most HF vision models return .last_hidden_state + # Skip the [CLS] token if present + features = outputs.last_hidden_state + if features.size(1) == self.num_patches + 1: + features = features[:, 1:, :] # remove CLS + return features + else: + # Fallback: simple conv patch embedding + # pixel_values: (B, 3, H, W) + x = self.patch_embed(pixel_values) # (B, D, H', W') + x = x.flatten(2).transpose(1, 2) # (B, num_patches, D) + return x + + +class MultiModalProjector(nn.Module): + """Maps vision encoder features to the LLM's embedding dimension. + + Supports two projector types: + - ``"linear"``: Single linear layer. + - ``"mlp2x"``: Two-layer MLP with GELU activation (LLaVA-style). + + Args: + vision_dim: Dimension of the vision encoder output. + text_dim: Dimension of the LLM's token embeddings (``config.n_embd``). + projector_type: ``"linear"`` or ``"mlp2x"``. + """ + + def __init__(self, vision_dim: int, text_dim: int, projector_type: str = "linear") -> None: + super().__init__() + self.projector_type = projector_type + + if projector_type == "linear": + self.proj = nn.Linear(vision_dim, text_dim, bias=True) + elif projector_type == "mlp2x": + self.proj = nn.Sequential( + nn.Linear(vision_dim, text_dim, bias=True), + nn.GELU(), + nn.Linear(text_dim, text_dim, bias=True), + ) + else: + raise ValueError(f"Unknown projector type: {projector_type!r}. Supported: 'linear', 'mlp2x'.") + + def forward(self, image_features: torch.Tensor) -> torch.Tensor: + """ + Args: + image_features: ``(B, num_patches, vision_dim)`` + + Returns: + Projected features of shape ``(B, num_patches, text_dim)``. + """ + return self.proj(image_features) + + +def merge_input_embeds( + text_embeds: torch.Tensor, + image_embeds: torch.Tensor, + image_token_id: int, + input_ids: torch.Tensor, +) -> torch.Tensor: + """Replace ```` placeholder embeddings with actual image embeddings. + + This function takes the standard text embedding output from ``wte`` and + splices in projected image patch embeddings at the positions where the + input contains the ``image_token_id`` placeholder token. + + Args: + text_embeds: ``(B, T, D)`` – embeddings from ``model.transformer.wte(idx)``. + image_embeds: ``(B, N_patches, D)`` – projected image patch embeddings. + image_token_id: The token ID used as the ```` placeholder. + input_ids: ``(B, T)`` – the original token IDs (needed to locate placeholders). + + Returns: + Merged embeddings ``(B, T, D)`` with image patches replacing placeholders. + + Raises: + ValueError: If the number of ```` placeholders doesn't match + the number of image patches. + """ + B, T, D = text_embeds.shape + N_patches = image_embeds.size(1) + + # Find positions of image placeholder tokens + image_mask = input_ids == image_token_id # (B, T) + + # Validate: each batch element should have exactly N_patches placeholders + counts = image_mask.sum(dim=1) # (B,) + if not (counts == N_patches).all(): + raise ValueError( + f"Expected {N_patches} placeholder tokens per sequence, but got counts: {counts.tolist()}" + ) + + # Clone so we don't modify the original + merged = text_embeds.clone() + + # For each batch element, scatter image embeddings + for b in range(B): + positions = image_mask[b].nonzero(as_tuple=True)[0] # (N_patches,) + merged[b, positions] = image_embeds[b] + + return merged + + +class ImagePreprocessor: + """Handles image loading, resizing, and normalization for VLMs. + + This class loads images from file paths or PIL Image objects, resizes + them to the expected input size, and normalizes pixel values. + + Args: + image_size: Target image size (both height and width). + mean: Per-channel normalization mean (default: ImageNet). + std: Per-channel normalization std (default: ImageNet). + """ + + # ImageNet defaults (used by CLIP, SigLIP, etc.) + IMAGENET_MEAN = (0.48145466, 0.4578275, 0.40821073) + IMAGENET_STD = (0.26862954, 0.26130258, 0.27577711) + + def __init__( + self, + image_size: int = 224, + mean: tuple[float, ...] = IMAGENET_MEAN, + std: tuple[float, ...] = IMAGENET_STD, + ) -> None: + self.image_size = image_size + self.mean = mean + self.std = std + + def __call__( + self, + image: str | Path | Any, + device: str | torch.device = "cpu", + ) -> torch.Tensor: + """Preprocess an image into a normalized tensor. + + Args: + image: A file path (str/Path) or a PIL Image object. + device: Target device for the output tensor. + + Returns: + ``(1, 3, image_size, image_size)`` tensor, normalized. + """ + try: + from PIL import Image as PILImage + except ImportError: + raise ImportError("Image preprocessing requires Pillow. Install it with: pip install Pillow") + + if isinstance(image, (str, Path)): + img = PILImage.open(image).convert("RGB") + else: + img = image.convert("RGB") + + # Resize with bicubic interpolation + img = img.resize((self.image_size, self.image_size), PILImage.BICUBIC) + + # Convert to tensor: (H, W, C) -> (C, H, W), scale to [0, 1] + import numpy as np + + pixel_values = torch.from_numpy(np.array(img, dtype=np.float32) / 255.0).permute(2, 0, 1) + + # Normalize + mean = torch.tensor(self.mean, dtype=pixel_values.dtype).view(3, 1, 1) + std = torch.tensor(self.std, dtype=pixel_values.dtype).view(3, 1, 1) + pixel_values = (pixel_values - mean) / std + + # Add batch dimension and move to device + return pixel_values.unsqueeze(0).to(device) diff --git a/tests/test_model.py b/tests/test_model.py index 8d0cf21d5e..a22996f3d4 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -1000,7 +1000,7 @@ def test_against_multimodal_gemma_3(model_name, device, dtype): ) ) - theirs_model = Gemma3ForConditionalGeneration(theirs_config).to(device) + theirs_model = Gemma3ForConditionalGeneration(theirs_config).to(device).eval() theirs_state_dict = theirs_model.state_dict() state_dict = {} diff --git a/tests/test_vision.py b/tests/test_vision.py new file mode 100644 index 0000000000..27a39cba20 --- /dev/null +++ b/tests/test_vision.py @@ -0,0 +1,323 @@ +# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file. + +"""Tests for litgpt.vision — multimodal support components.""" + +import pytest +import torch +import torch.nn as nn + +from litgpt.config import Config +from litgpt.vision import ( + ImagePreprocessor, + MultiModalProjector, + VisionEncoder, + merge_input_embeds, +) + + +# --------------------------------------------------------------------------- +# Helper to build a minimal multimodal Config +# --------------------------------------------------------------------------- +def _mm_config(**overrides): + """Build a minimal Config suitable for multimodal tests.""" + defaults = dict( + name="test-mm", + hf_config={"name": "test-mm", "org": "test"}, + block_size=128, + vocab_size=512, + n_layer=2, + n_head=4, + n_embd=64, + padding_multiple=64, + # Vision fields + vision_feature_dim=32, + vision_start_token_id=100, + vision_patch_size=14, + vision_image_size=28, # small to keep tests fast + mm_projector_type="linear", + ) + defaults.update(overrides) + return Config(**defaults) + + +def _text_only_config(**overrides): + """Build a minimal text-only Config (no vision fields).""" + defaults = dict( + name="test-text", + hf_config={"name": "test-text", "org": "test"}, + block_size=128, + vocab_size=512, + n_layer=2, + n_head=4, + n_embd=64, + padding_multiple=64, + ) + defaults.update(overrides) + return Config(**defaults) + + +# ===== Config Tests ===== +class TestConfigMultimodal: + def test_is_multimodal_true(self): + config = _mm_config() + assert config.is_multimodal is True + + def test_is_multimodal_false(self): + config = _text_only_config() + assert config.is_multimodal is False + + def test_vision_fields_default_to_none(self): + config = _text_only_config() + assert config.vision_feature_dim is None + assert config.vision_start_token_id is None + assert config.vision_patch_size is None + assert config.vision_image_size is None + assert config.mm_projector_type is None + assert config.vision_model_name is None + + +# ===== VisionEncoder Tests ===== +class TestVisionEncoder: + def test_fallback_forward_shape(self): + config = _mm_config() + encoder = VisionEncoder(config) + # image_size=28, patch_size=14 -> 2x2 = 4 patches + pixel_values = torch.randn(1, 3, 28, 28) + features = encoder(pixel_values) + assert features.shape == (1, 4, 32) # (B, num_patches, vision_feature_dim) + + def test_fallback_batched(self): + config = _mm_config() + encoder = VisionEncoder(config) + pixel_values = torch.randn(3, 3, 28, 28) + features = encoder(pixel_values) + assert features.shape == (3, 4, 32) + + def test_requires_vision_feature_dim(self): + config = _text_only_config() + with pytest.raises(ValueError, match="vision_feature_dim"): + VisionEncoder(config) + + def test_num_patches_property(self): + config = _mm_config() + encoder = VisionEncoder(config) + assert encoder.num_patches == 4 # (28/14)^2 + + +# ===== MultiModalProjector Tests ===== +class TestMultiModalProjector: + def test_linear_projector_shape(self): + proj = MultiModalProjector(vision_dim=32, text_dim=64, projector_type="linear") + x = torch.randn(2, 4, 32) + out = proj(x) + assert out.shape == (2, 4, 64) + + def test_mlp2x_projector_shape(self): + proj = MultiModalProjector(vision_dim=32, text_dim=64, projector_type="mlp2x") + x = torch.randn(2, 4, 32) + out = proj(x) + assert out.shape == (2, 4, 64) + + def test_invalid_projector_type(self): + with pytest.raises(ValueError, match="Unknown projector type"): + MultiModalProjector(vision_dim=32, text_dim=64, projector_type="bad") + + def test_linear_is_single_layer(self): + proj = MultiModalProjector(vision_dim=32, text_dim=64, projector_type="linear") + assert isinstance(proj.proj, nn.Linear) + + def test_mlp2x_is_sequential(self): + proj = MultiModalProjector(vision_dim=32, text_dim=64, projector_type="mlp2x") + assert isinstance(proj.proj, nn.Sequential) + assert len(proj.proj) == 3 # Linear + GELU + Linear + + +# ===== merge_input_embeds Tests ===== +class TestMergeInputEmbeds: + def test_basic_merge(self): + B, T, D = 1, 10, 64 + N_patches = 4 + image_token_id = 999 + + # Create input_ids with 4 image placeholders at positions 2,3,4,5 + input_ids = torch.arange(T).unsqueeze(0) # (1, 10) + input_ids[0, 2:6] = image_token_id + + text_embeds = torch.zeros(B, T, D) + image_embeds = torch.ones(B, N_patches, D) + + merged = merge_input_embeds(text_embeds, image_embeds, image_token_id, input_ids) + assert merged.shape == (B, T, D) + + # Positions 2-5 should have image embeddings (all ones) + assert torch.allclose(merged[0, 2:6], torch.ones(N_patches, D)) + # Other positions should still be zeros + assert torch.allclose(merged[0, :2], torch.zeros(2, D)) + assert torch.allclose(merged[0, 6:], torch.zeros(4, D)) + + def test_batch_merge(self): + B, T, D = 2, 8, 32 + N_patches = 2 + image_token_id = 999 + + input_ids = torch.zeros(B, T, dtype=torch.long) + input_ids[0, 1:3] = image_token_id + input_ids[1, 5:7] = image_token_id + + text_embeds = torch.zeros(B, T, D) + image_embeds = torch.ones(B, N_patches, D) * 2.0 + + merged = merge_input_embeds(text_embeds, image_embeds, image_token_id, input_ids) + assert merged.shape == (B, T, D) + + # Check first batch element + assert merged[0, 1, 0].item() == 2.0 + assert merged[0, 2, 0].item() == 2.0 + assert merged[0, 0, 0].item() == 0.0 + + # Check second batch element + assert merged[1, 5, 0].item() == 2.0 + assert merged[1, 6, 0].item() == 2.0 + assert merged[1, 0, 0].item() == 0.0 + + def test_wrong_placeholder_count_raises(self): + B, T, D = 1, 8, 32 + N_patches = 4 + image_token_id = 999 + + input_ids = torch.zeros(B, T, dtype=torch.long) + input_ids[0, 0:2] = image_token_id # Only 2, not 4 + + text_embeds = torch.zeros(B, T, D) + image_embeds = torch.ones(B, N_patches, D) + + with pytest.raises(ValueError, match="Expected 4"): + merge_input_embeds(text_embeds, image_embeds, image_token_id, input_ids) + + def test_does_not_modify_original(self): + B, T, D = 1, 6, 16 + N_patches = 2 + image_token_id = 999 + + input_ids = torch.zeros(B, T, dtype=torch.long) + input_ids[0, 1:3] = image_token_id + + text_embeds = torch.zeros(B, T, D) + image_embeds = torch.ones(B, N_patches, D) + + original = text_embeds.clone() + merge_input_embeds(text_embeds, image_embeds, image_token_id, input_ids) + assert torch.allclose(text_embeds, original), "Original embeddings should not be modified" + + +# ===== GPT Model Integration Tests ===== +class TestGPTMultimodal: + def test_text_only_model_unchanged(self): + """Text-only models should work exactly as before.""" + config = _text_only_config() + from litgpt.model import GPT + + model = GPT(config) + assert model.vision_encoder is None + assert model.mm_projector is None + + # Forward should work without pixel_values + idx = torch.randint(0, config.padded_vocab_size, (1, 5)) + output = model(idx) + assert output.shape == (1, 5, config.padded_vocab_size) + + def test_multimodal_model_has_vision_components(self): + config = _mm_config() + from litgpt.model import GPT + + model = GPT(config) + assert model.vision_encoder is not None + assert model.mm_projector is not None + + def test_multimodal_forward_with_pixel_values(self): + """Test that forward works when pixel_values are provided.""" + config = _mm_config() + from litgpt.model import GPT + + model = GPT(config) + model.eval() + + # Create input with image placeholders + N_patches = 4 # (28/14)^2 + T = 10 + idx = torch.randint(0, config.padded_vocab_size, (1, T)) + # Place image tokens at positions 2-5 + idx[0, 2:6] = config.vision_start_token_id + + pixel_values = torch.randn(1, 3, 28, 28) + + with torch.no_grad(): + output = model(idx, pixel_values=pixel_values) + + assert output.shape == (1, T, config.padded_vocab_size) + + def test_multimodal_forward_without_pixel_values(self): + """Multimodal model should still work for text-only inference.""" + config = _mm_config() + from litgpt.model import GPT + + model = GPT(config) + model.eval() + + idx = torch.randint(0, config.padded_vocab_size, (1, 5)) + + with torch.no_grad(): + output = model(idx) + + assert output.shape == (1, 5, config.padded_vocab_size) + + +# ===== ImagePreprocessor Tests ===== +class TestImagePreprocessor: + def test_output_shape(self, tmp_path): + """Test that preprocessor produces correct tensor shape.""" + # Create a dummy image file + try: + from PIL import Image + except ImportError: + pytest.skip("Pillow not installed") + + img = Image.new("RGB", (100, 100), color=(128, 64, 32)) + img_path = tmp_path / "test.jpg" + img.save(str(img_path)) + + preprocessor = ImagePreprocessor(image_size=28) + result = preprocessor(str(img_path)) + + assert result.shape == (1, 3, 28, 28) + assert result.dtype == torch.float32 + + def test_pil_image_input(self): + """Test that preprocessor accepts PIL Image objects.""" + try: + from PIL import Image + except ImportError: + pytest.skip("Pillow not installed") + + img = Image.new("RGB", (50, 50), color=(255, 0, 0)) + preprocessor = ImagePreprocessor(image_size=14) + result = preprocessor(img) + assert result.shape == (1, 3, 14, 14) + + def test_normalization(self, tmp_path): + """Test that output is normalized (not in [0, 255] range).""" + try: + from PIL import Image + except ImportError: + pytest.skip("Pillow not installed") + + img = Image.new("RGB", (28, 28), color=(128, 128, 128)) + img_path = tmp_path / "test.png" + img.save(str(img_path)) + + preprocessor = ImagePreprocessor(image_size=28) + result = preprocessor(str(img_path)) + + # After normalization, values should not be in raw [0, 255] range + assert result.max().item() < 10.0 + assert result.min().item() > -10.0