lightvl 0.0.3__py3-none-any.whl

lightvl/__init__.py

@@ -0,0 +1 @@
+#lightvl

lightvl/backends/__init__.py

@@ -0,0 +1,49 @@
+# Inference backends
+
+from .base_backend import BaseBackend
+from .data_utils import (
+    load_dataset,
+    load_image_as_pil,
+    preprocess_image,
+    process_image_url,
+    get_sample_param_from_config,
+    get_model_type,
+    parse_max_img_size,
+)
+from .transformers_backend import (
+    TransformersBackend,
+    run_transformers_inference,
+)
+from .vllm_backend import (
+    VLLMBackend,
+    run_vllm_inference,
+    run_models_on_dataset,
+    cleanup_vllm_processes,
+    start_vllm_serve,
+    stop_vllm_serve,
+    wait_for_vllm_ready,
+)
+
+__all__ = [
+    # Base
+    "BaseBackend",
+    # Data utils
+    "load_dataset",
+    "load_image_as_pil",
+    "preprocess_image",
+    "process_image_url",
+    "get_sample_param_from_config",
+    "get_model_type",
+    "parse_max_img_size",
+    # Transformers
+    "TransformersBackend",
+    "run_transformers_inference",
+    # vLLM
+    "VLLMBackend",
+    "run_vllm_inference",
+    "run_models_on_dataset",
+    "cleanup_vllm_processes",
+    "start_vllm_serve",
+    "stop_vllm_serve",
+    "wait_for_vllm_ready",
+]

lightvl/backends/base_backend.py

@@ -0,0 +1,156 @@
+"""Abstract base backend for Vision-Language model inference."""
+
+import time
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+from loguru import logger
+from PIL import Image as PILImage
+from tqdm import tqdm
+
+from lightvl.backends.data_utils import (
+    load_dataset,
+    load_image_as_pil,
+    get_sample_param_from_config,
+    parse_max_img_size,
+)
+
+
+class BaseBackend(ABC):
+    """
+    Abstract base class for inference backends.
+
+    Subclasses implement:
+      - load_model() -> (model, processor/tokenizer)
+      - infer_sample() -> (output_text, prompt_tokens, completion_tokens)
+      - cleanup() -> release GPU resources
+
+    The shared run_inference() handles dataset loading, progress bar,
+    per-sample image loading + timing, error handling, and result formatting.
+    """
+
+    def __init__(self, model_path: str, config: Optional[Dict[str, Any]] = None):
+        self.model_path = model_path
+        self.config = config or {}
+        self.model = None
+        self.processor = None
+
+        # Parsed config values
+        self.sampling_params = get_sample_param_from_config(self.config)
+        self.sys_prompt = self.config.get("sys_prompt", "")
+        self.usr_prompt = self.config.get("usr_prompt", "")
+        self.max_img_size = parse_max_img_size(self.config)
+
+    @abstractmethod
+    def load_model(self) -> None:
+        """
+        Load model and processor/tokenizer.
+        Sets self.model and self.processor.
+        Must be called before run_inference().
+        """
+        ...
+
+    @abstractmethod
+    def infer_sample(
+        self,
+        pil_image: Optional[PILImage.Image],
+        prompt: str,
+    ) -> Tuple[str, int, int]:
+        """
+        Run inference on a single sample.
+        Returns (output_text, prompt_tokens, completion_tokens).
+        """
+        ...
+
+    @abstractmethod
+    def get_cmd_str(self) -> str:
+        """Return a human-readable command string describing how the model was loaded."""
+        ...
+  
+    def run_inference(self, dataset_path: str) -> Tuple[List[Dict[str, Any]], str]:
+        """
+        Run inference on all samples in the dataset.
+
+        Shared implementation handles:
+          - Dataset loading
+          - Per-sample image loading (excluded from latency)
+          - Inference timing
+          - Progress bar display
+          - Error handling and result formatting
+
+        :return: (results_list, command_str)
+        """
+        data_list = load_dataset(dataset_path)
+        logger.info("Running inference on {} samples", len(data_list))
+
+        results = []
+        time.sleep(0.1)  # allow GPU to settle after model load
+
+        for idx, item in enumerate(
+            tqdm(data_list, desc=f"{self.__class__.__name__} inference", unit="sample")
+        ):
+            prompt = item.get("prompt", "") or self.usr_prompt
+            image_url = item.get("image_url", "")
+            ground_truth = item.get("ground_truth", "")
+
+            start_time = None
+            try:
+                # Image loading is NOT counted in inference latency
+                pil_image = None
+                if image_url:
+                    pil_image = load_image_as_pil(image_url, self.max_img_size)
+
+                start_time = time.time()
+
+                output_text, prompt_tokens, completion_tokens = self.infer_sample(
+                    pil_image, prompt
+                )
+
+                # Read first-token logits if backend provides them
+                first_token_logits = getattr(self, '_last_first_token_logits', None)
+
+                latency_ms = (time.time() - start_time) * 1000
+                ttft_ms = 0.0  # Not measurable without streaming
+                tpot_ms = 0.0  # Not measurable without streaming
+
+                result = {
+                    "index": idx,
+                    "input": {
+                        "image_url": image_url,
+                        "prompt": prompt,
+                        "ground_truth": ground_truth,
+                    },
+                    "output": output_text,
+                    "latency_ms": latency_ms,
+                    "ttft_ms": ttft_ms,
+                    "tpot_ms": tpot_ms,
+                    "prompt_tokens": prompt_tokens,
+                    "completion_tokens": completion_tokens,
+                    "first_token_logits": first_token_logits,
+                }
+            except Exception as e:
+                latency_ms = (time.time() - start_time) * 1000 if start_time is not None else 0
+                logger.error("Inference failed for sample {}: {}", idx, str(e))
+                result = {
+                    "index": idx,
+                    "input": {
+                        "image_url": image_url,
+                        "prompt": prompt,
+                        "ground_truth": ground_truth,
+                    },
+                    "output": None,
+                    "latency_ms": latency_ms,
+                    "ttft_ms": 0.0,
+                    "tpot_ms": 0.0,
+                    "prompt_tokens": 0,
+                    "completion_tokens": 0,
+                    "error": str(e),
+                }
+
+            results.append(result)
+
+        cmd_str = self.get_cmd_str()
+        logger.info("{} inference completed: {} results", self.__class__.__name__, len(results))
+
+        return results, cmd_str

lightvl/backends/data_utils.py

@@ -0,0 +1,206 @@
+"""Shared data loading, image preprocessing, and config utilities for all backends."""
+
+import ast
+import csv
+import io
+import math
+import os
+import base64
+import json
+from typing import Any, Dict, List, Tuple
+
+import requests
+from PIL import Image
+from loguru import logger
+
+
+# ===================== Dataset Loading =====================
+
+def load_dataset(dataset_path: str) -> List[Dict[str, Any]]:
+    """
+    Load dataset from CSV file with unified format for both VL and embedding tasks.
+    CSV format: prompt, image_url, ground_truth
+    :param dataset_path: Dataset path (CSV file)
+    :return: Data list with unified format: {"image_url": str, "prompt": str, "ground_truth": str}
+    """
+    data_list = []
+    logger.info("Loading dataset file: {}", dataset_path)
+
+    if dataset_path.endswith(".csv"):
+        with open(dataset_path, "r", encoding="utf-8") as f:
+            reader = csv.DictReader(f)
+            for row in reader:
+                item = {
+                    "image_url": row.get("image_url", ""),
+                    "prompt": row.get("prompt", ""),
+                    "ground_truth": row.get("ground_truth", ""),
+                }
+                data_list.append(item)
+
+    # Fill empty prompts with the first row's prompt
+    if data_list:
+        first_prompt = data_list[0].get("prompt", "")
+        if first_prompt:
+            for item in data_list:
+                if not item.get("prompt", ""):
+                    item["prompt"] = first_prompt
+
+    logger.info("Loaded {} data entries", len(data_list))
+    return data_list
+
+
+# ===================== Image Preprocessing =====================
+
+def preprocess_image(
+        image: Image.Image,
+        target_size: Tuple[int, int],
+        min_size: Tuple[int, int] = (32, 32),
+) -> Image.Image:
+    """
+    Pre-process an image by scaling it to fit within pixel bounds.
+
+    If the image exceeds ``target_size[0] * target_size[1]`` pixels, scale it
+    down proportionally.  If it is smaller than ``min_size[0] * min_size[1]``
+    pixels, scale it up proportionally.  Always returns an RGB image.
+
+    :param image: PIL Image
+    :param target_size: Maximum pixel dimensions (width, height).  The product
+                        defines **image_max_pixels**.
+    :param min_size: Minimum pixel dimensions (width, height).  The product
+                     defines **image_min_pixels**.
+    :return: Pre-processed PIL Image in RGB mode
+    """
+    image_max_pixels = target_size[0] * target_size[1]
+    image_min_pixels = min_size[0] * min_size[1]
+
+    img_w = image.width
+    img_h = image.height
+    img_pixels = img_w * img_h
+
+    if img_pixels > image_max_pixels:
+        resize_factor = math.sqrt(image_max_pixels / img_pixels)
+        new_w, new_h = int(img_w * resize_factor), int(img_h * resize_factor)
+        image = image.resize((new_w, new_h), Image.NEAREST)
+
+    if img_pixels < image_min_pixels:
+        resize_factor = math.sqrt(image_min_pixels / img_pixels)
+        new_w, new_h = int(img_w * resize_factor), int(img_h * resize_factor)
+        image = image.resize((new_w, new_h), Image.NEAREST)
+
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    return image
+
+
+def load_image_as_pil(image_url: str, target_size: Tuple[int, int] = None) -> Image.Image:
+    """
+    Download image or load local image, resize and return PIL Image.
+
+    Uses :func:`preprocess_image` for pixel-bound-aware resizing.
+
+    :param image_url: Image URL or local file path
+    :param target_size: Target size tuple (width, height), e.g. (512, 512).
+                        Used to compute max_pixels = w * h.
+    :return: PIL Image in RGB mode
+    :raises ValueError: If image source is invalid
+    """
+    image_bytes = None
+
+    if os.path.isfile(image_url):
+        logger.debug("Loading local image: {}", image_url)
+        with open(image_url, "rb") as f:
+            image_bytes = f.read()
+    elif image_url.startswith(("http://", "https://")):
+        response = requests.get(image_url, timeout=10)
+        response.raise_for_status()
+        image_bytes = response.content
+    else:
+        raise ValueError(f"Invalid image source: not a valid URL or local file path: {image_url}")
+
+    image = Image.open(io.BytesIO(image_bytes))
+
+    if target_size is not None:
+        image = preprocess_image(image, target_size)
+    elif image.mode != "RGB":
+        image = image.convert("RGB")
+
+    return image
+
+
+def process_image_url(image_url: str, target_size: Tuple[int, int] = None) -> str:
+    """
+    Download image or load local image, resize and convert to base64 format.
+    :param image_url: Image URL or local file path
+    :param target_size: Target size, default 512x512
+    :return: Base64 encoded image data (usable by vLLM)
+    """
+    try:
+        image = load_image_as_pil(image_url, target_size)
+
+        buffer = io.BytesIO()
+        image.save(buffer, format="PNG")
+        image_bytes = buffer.getvalue()
+        base64_str = base64.b64encode(image_bytes).decode("utf-8")
+
+        return f"data:image/png;base64,{base64_str}"
+
+    except Exception as e:
+        logger.error("Image processing failed: {}", str(e))
+        return ""
+
+
+# ===================== Config Helpers =====================
+
+def get_sample_param_from_config(config: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Extract sampling parameters from config dict.
+    :param config: Configuration dictionary
+    :return: Sampling parameters dict
+    """
+    return {
+        "temperature": config.get("temperature", 0.0),
+        "top_p": config.get("top_p", 1.0),
+        "n": config.get("n", 1),
+        "top_k": config.get("top_k", 1),
+        "repetition_penalty": config.get("repetition_penalty", 1.0),
+        "max_tokens": config.get("max_tokens", 512),
+        "stop": config.get("stop", []),
+        "skip_special_tokens": config.get("skip_special_tokens", False),
+        "stream": config.get("stream", False),
+        "enable_thinking": config.get("enable_thinking", False),
+    }
+
+
+def get_model_type(model_path: str) -> str:
+    """Read model_type from config.json in the model directory."""
+    config_path = os.path.join(model_path, "config.json")
+    if not os.path.exists(config_path):
+        raise FileNotFoundError(f"config.json not found in {model_path}")
+
+    with open(config_path, "r", encoding="utf-8") as f:
+        config = json.load(f)
+
+    model_type = config.get("model_type", "")
+    if not model_type:
+        raise ValueError(f"model_type not found in config.json for {model_path}")
+
+    return model_type
+
+
+def parse_max_img_size(config: Dict[str, Any]) -> Tuple[int, int]:
+    """
+    Parse max_img_size from config, supporting both string and tuple formats.
+    :param config: Configuration dictionary
+    :return: (width, height) tuple
+    """
+    max_img_size = config.get("max_img_size", None)
+    if max_img_size:
+        try:
+            max_img_size = ast.literal_eval(max_img_size)
+        except Exception:
+            max_img_size = (512, 512)
+    else:
+        max_img_size = (512, 512)
+
+    return max_img_size