{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Query Vision Language Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## Querying Qwen-VL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()  # Run this first.\n",
    "\n",
    "model_path = \"Qwen/Qwen2.5-VL-3B-Instruct\"\n",
    "chat_template = \"qwen2-vl\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|im_start|>system\n",
      "You are a helpful assistant.<|im_end|>\n",
      "<|im_start|>user\n",
      "What's shown here: <|vision_start|><|image_pad|><|vision_end|>?<|im_end|>\n",
      "<|im_start|>assistant\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 55528 characters]",
      "image/png": "[encoded data, 443712 characters]",
      "text/plain": [
       "<PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=570x380>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Lets create a prompt.\n",
    "\n",
    "from io import BytesIO\n",
    "import requests\n",
    "from PIL import Image\n",
    "\n",
    "from sglang.srt.conversation import chat_templates\n",
    "\n",
    "image = Image.open(\n",
    "    BytesIO(\n",
    "        requests.get(\n",
    "            \"https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true\"\n",
    "        ).content\n",
    "    )\n",
    ")\n",
    "\n",
    "conv = chat_templates[chat_template].copy()\n",
    "conv.append_message(conv.roles[0], f\"What's shown here: {conv.image_token}?\")\n",
    "conv.append_message(conv.roles[1], \"\")\n",
    "conv.image_data = [image]\n",
    "\n",
    "print(conv.get_prompt())\n",
    "image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "### Query via the offline Engine API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You have video processor config saved in `preprocessor.json` file which is deprecated. Video processor configs should be saved in their own `video_preprocessor.json` file. You can rename the file or load and save the processor back which renames it automatically. Loading from `preprocessor.json` will be removed in v5.0.\n",
      "You have video processor config saved in `preprocessor.json` file which is deprecated. Video processor configs should be saved in their own `video_preprocessor.json` file. You can rename the file or load and save the processor back which renames it automatically. Loading from `preprocessor.json` will be removed in v5.0.\n",
      "Loading safetensors checkpoint shards:   0% Completed | 0/2 [00:00<?, ?it/s]\n",
      "Loading safetensors checkpoint shards:  50% Completed | 1/2 [00:03<00:03,  3.13s/it]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 2/2 [00:06<00:00,  3.27s/it]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 2/2 [00:06<00:00,  3.25s/it]\n",
      "\n",
      "Capturing batches (bs=1 avail_mem=21.63 GB): 100%|██████████| 35/35 [00:10<00:00,  3.19it/s]  \n"
     ]
    }
   ],
   "source": [
    "from sglang import Engine\n",
    "\n",
    "llm = Engine(\n",
    "    model_path=model_path, chat_template=chat_template, mem_fraction_static=0.8\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "In the picture, a person in a yellow shirt is hanging laundry on a clothesline attached to the back of a yellow taxi in an urban setting. There are city streets, buildings, and traffic lights visible in the background. The scene appears to be incongruous and amusing, as it shows an unusual and somewhat chaotic activity happening in a busy city environment.\n"
     ]
    }
   ],
   "source": [
    "out = llm.generate(prompt=conv.get_prompt(), image_data=[image])\n",
    "print(out[\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "### Query via the offline Engine API, but send precomputed embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7c94dead4660409c9acfac1f3461d7d9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Compute the image embeddings using Huggingface.\n",
    "\n",
    "from transformers import AutoProcessor\n",
    "from transformers import Qwen2_5_VLForConditionalGeneration\n",
    "\n",
    "processor = AutoProcessor.from_pretrained(model_path, use_fast=True)\n",
    "vision = (\n",
    "    Qwen2_5_VLForConditionalGeneration.from_pretrained(model_path).eval().visual.cuda()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The image shows a scene with two yellow taxis in an urban setting. The taxi on the left has a red light on top, indicating that it may be waiting or preparing to drive. The other taxi, which is facing left, has its hatch open with some clothing or fabric hanging out. The background features high-rise buildings and city streets, suggesting this is taking place in a downtown area of a city. The presence of multiple flags on flagpoles indicates that there might be some celebration or event within the vicinity.\n"
     ]
    }
   ],
   "source": [
    "processed_prompt = processor(\n",
    "    images=[image], text=conv.get_prompt(), return_tensors=\"pt\"\n",
    ")\n",
    "input_ids = processed_prompt[\"input_ids\"][0].detach().cpu().tolist()\n",
    "precomputed_embeddings = vision(\n",
    "    processed_prompt[\"pixel_values\"].cuda(), processed_prompt[\"image_grid_thw\"].cuda()\n",
    ")\n",
    "\n",
    "mm_item = dict(\n",
    "    modality=\"IMAGE\",\n",
    "    image_grid_thw=processed_prompt[\"image_grid_thw\"],\n",
    "    precomputed_embeddings=precomputed_embeddings,\n",
    ")\n",
    "out = llm.generate(input_ids=input_ids, image_data=[mm_item])\n",
    "print(out[\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "## Querying Llama 4 (Vision)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()  # Run this first.\n",
    "\n",
    "model_path = \"meta-llama/Llama-4-Scout-17B-16E-Instruct\"\n",
    "chat_template = \"llama-4\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|header_start|>user<|header_end|>\n",
      "\n",
      "What's shown here: <|image|>?<|eot|><|header_start|>assistant<|header_end|>\n",
      "\n",
      "\n",
      "Image size: (570, 380)\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 55528 characters]",
      "image/png": "[encoded data, 443712 characters]",
      "text/plain": [
       "<PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=570x380>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Lets create a prompt.\n",
    "\n",
    "from io import BytesIO\n",
    "import requests\n",
    "from PIL import Image\n",
    "\n",
    "from sglang.srt.conversation import chat_templates\n",
    "\n",
    "image = Image.open(\n",
    "    BytesIO(\n",
    "        requests.get(\n",
    "            \"https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true\"\n",
    "        ).content\n",
    "    )\n",
    ")\n",
    "\n",
    "conv = chat_templates[chat_template].copy()\n",
    "conv.append_message(conv.roles[0], f\"What's shown here: {conv.image_token}?\")\n",
    "conv.append_message(conv.roles[1], \"\")\n",
    "conv.image_data = [image]\n",
    "\n",
    "print(conv.get_prompt())\n",
    "print(f\"Image size: {image.size}\")\n",
    "\n",
    "image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "### Query via the offline Engine API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading safetensors checkpoint shards:   0% Completed | 0/50 [00:00<?, ?it/s]\n",
      "Loading safetensors checkpoint shards:   2% Completed | 1/50 [00:22<18:10, 22.26s/it]\n",
      "Loading safetensors checkpoint shards:   4% Completed | 2/50 [00:44<17:44, 22.17s/it]\n",
      "Loading safetensors checkpoint shards:   6% Completed | 3/50 [01:06<17:24, 22.22s/it]\n",
      "Loading safetensors checkpoint shards:   8% Completed | 4/50 [01:28<16:55, 22.07s/it]\n",
      "Loading safetensors checkpoint shards:  10% Completed | 5/50 [01:50<16:28, 21.96s/it]\n",
      "Loading safetensors checkpoint shards:  12% Completed | 6/50 [02:11<15:59, 21.80s/it]\n",
      "Loading safetensors checkpoint shards:  14% Completed | 7/50 [02:34<15:52, 22.14s/it]\n",
      "Loading safetensors checkpoint shards:  16% Completed | 8/50 [02:54<15:05, 21.57s/it]\n",
      "Loading safetensors checkpoint shards:  18% Completed | 9/50 [03:17<14:51, 21.74s/it]\n",
      "Loading safetensors checkpoint shards:  20% Completed | 10/50 [03:29<12:31, 18.79s/it]\n",
      "Loading safetensors checkpoint shards:  22% Completed | 11/50 [03:32<09:10, 14.13s/it]\n",
      "Loading safetensors checkpoint shards:  24% Completed | 12/50 [03:36<06:53, 10.89s/it]\n",
      "Loading safetensors checkpoint shards:  26% Completed | 13/50 [03:39<05:19,  8.65s/it]\n",
      "Loading safetensors checkpoint shards:  28% Completed | 14/50 [03:43<04:15,  7.09s/it]\n",
      "Loading safetensors checkpoint shards:  30% Completed | 15/50 [03:46<03:29,  6.00s/it]\n",
      "Loading safetensors checkpoint shards:  32% Completed | 16/50 [03:50<02:57,  5.23s/it]\n",
      "Loading safetensors checkpoint shards:  34% Completed | 17/50 [03:53<02:35,  4.73s/it]\n",
      "Loading safetensors checkpoint shards:  36% Completed | 18/50 [03:57<02:18,  4.33s/it]\n",
      "Loading safetensors checkpoint shards:  38% Completed | 19/50 [04:00<02:06,  4.09s/it]\n",
      "Loading safetensors checkpoint shards:  40% Completed | 20/50 [04:04<01:56,  3.87s/it]\n",
      "Loading safetensors checkpoint shards:  42% Completed | 21/50 [04:07<01:48,  3.74s/it]\n",
      "Loading safetensors checkpoint shards:  44% Completed | 22/50 [04:11<01:43,  3.71s/it]\n",
      "Loading safetensors checkpoint shards:  46% Completed | 23/50 [04:14<01:37,  3.63s/it]\n",
      "Loading safetensors checkpoint shards:  48% Completed | 24/50 [04:18<01:33,  3.60s/it]\n",
      "Loading safetensors checkpoint shards:  50% Completed | 25/50 [04:21<01:26,  3.45s/it]\n",
      "Loading safetensors checkpoint shards:  52% Completed | 26/50 [04:21<01:02,  2.61s/it]\n",
      "Loading safetensors checkpoint shards:  54% Completed | 27/50 [04:25<01:06,  2.91s/it]\n",
      "Loading safetensors checkpoint shards:  56% Completed | 28/50 [04:28<01:07,  3.09s/it]\n",
      "Loading safetensors checkpoint shards:  58% Completed | 29/50 [04:32<01:07,  3.20s/it]\n",
      "Loading safetensors checkpoint shards:  60% Completed | 30/50 [04:35<01:05,  3.25s/it]\n",
      "Loading safetensors checkpoint shards:  62% Completed | 31/50 [04:39<01:02,  3.30s/it]\n",
      "Loading safetensors checkpoint shards:  64% Completed | 32/50 [04:42<01:00,  3.37s/it]\n",
      "Loading safetensors checkpoint shards:  66% Completed | 33/50 [04:46<00:58,  3.45s/it]\n",
      "Loading safetensors checkpoint shards:  68% Completed | 34/50 [04:49<00:55,  3.45s/it]\n",
      "Loading safetensors checkpoint shards:  70% Completed | 35/50 [04:53<00:51,  3.45s/it]\n",
      "Loading safetensors checkpoint shards:  72% Completed | 36/50 [04:56<00:48,  3.46s/it]\n",
      "Loading safetensors checkpoint shards:  74% Completed | 37/50 [05:00<00:44,  3.45s/it]\n",
      "Loading safetensors checkpoint shards:  76% Completed | 38/50 [05:03<00:41,  3.45s/it]\n",
      "Loading safetensors checkpoint shards:  78% Completed | 39/50 [05:07<00:38,  3.50s/it]\n",
      "Loading safetensors checkpoint shards:  80% Completed | 40/50 [05:10<00:34,  3.49s/it]\n",
      "Loading safetensors checkpoint shards:  82% Completed | 41/50 [05:14<00:31,  3.49s/it]\n",
      "Loading safetensors checkpoint shards:  84% Completed | 42/50 [05:17<00:27,  3.47s/it]\n",
      "Loading safetensors checkpoint shards:  86% Completed | 43/50 [05:20<00:24,  3.43s/it]\n",
      "Loading safetensors checkpoint shards:  88% Completed | 44/50 [05:24<00:20,  3.46s/it]\n",
      "Loading safetensors checkpoint shards:  90% Completed | 45/50 [05:27<00:17,  3.44s/it]\n",
      "Loading safetensors checkpoint shards:  92% Completed | 46/50 [05:31<00:13,  3.44s/it]\n",
      "Loading safetensors checkpoint shards:  94% Completed | 47/50 [05:34<00:10,  3.43s/it]\n",
      "Loading safetensors checkpoint shards:  96% Completed | 48/50 [05:38<00:06,  3.43s/it]\n",
      "Loading safetensors checkpoint shards:  98% Completed | 49/50 [05:41<00:03,  3.45s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setting sliding_window_size to be attention_chunk_size: 8192Setting sliding_window_size to be attention_chunk_size: 8192\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading safetensors checkpoint shards: 100% Completed | 50/50 [05:44<00:00,  3.43s/it]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 50/50 [05:44<00:00,  6.90s/it]\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setting sliding_window_size to be attention_chunk_size: 8192\n",
      "Setting sliding_window_size to be attention_chunk_size: 8192\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Capturing batches (bs=1 avail_mem=21.53 GB): 100%|██████████| 35/35 [00:15<00:00,  2.25it/s]  \n"
     ]
    }
   ],
   "source": [
    "from sglang.test.test_utils import is_in_ci\n",
    "\n",
    "if not is_in_ci():\n",
    "    from sglang import Engine\n",
    "\n",
    "    llm = Engine(\n",
    "        model_path=model_path,\n",
    "        trust_remote_code=True,\n",
    "        enable_multimodal=True,\n",
    "        mem_fraction_static=0.8,\n",
    "        tp_size=4,\n",
    "        attention_backend=\"fa3\",\n",
    "        context_length=65536,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The image depicts a man ironing clothing on the back of a yellow SUV in a city street, with another yellow taxi passing by. The man is wearing a yellow shirt and appears to be ironing a blue shirt on a makeshift ironing board set up behind the SUV. The scene suggests that the man may be a street vendor or someone who is trying to make a living by providing ironing services to people on the go.\n"
     ]
    }
   ],
   "source": [
    "if not is_in_ci():\n",
    "    out = llm.generate(prompt=conv.get_prompt(), image_data=[image])\n",
    "    print(out[\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "### Query via the offline Engine API, but send precomputed embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0eae2e36d07d42b89bc4b5ac7d62f226",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "if not is_in_ci():\n",
    "    # Compute the image embeddings using Huggingface.\n",
    "\n",
    "    from transformers import AutoProcessor\n",
    "    from transformers import Llama4ForConditionalGeneration\n",
    "\n",
    "    processor = AutoProcessor.from_pretrained(model_path, use_fast=True)\n",
    "    model = Llama4ForConditionalGeneration.from_pretrained(\n",
    "        model_path, torch_dtype=\"auto\"\n",
    "    ).eval()\n",
    "    vision = model.vision_model.cuda()\n",
    "    multi_modal_projector = model.multi_modal_projector.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "processed_prompt[\"pixel_values\"].shape=torch.Size([5, 3, 336, 336])\n",
      "The image depicts a man ironing on a makeshift ironing board set up on the back of a yellow SUV, in the middle of a busy street. The man is wearing a yellow shirt and appears to be ironing a blue shirt. In the background, there are other yellow taxis and tall buildings, suggesting that the scene is set in a city, likely New York City. The overall scene is one of a person going about their daily activities in a busy urban environment.\n"
     ]
    }
   ],
   "source": [
    "if not is_in_ci():\n",
    "    processed_prompt = processor(\n",
    "        images=[image], text=conv.get_prompt(), return_tensors=\"pt\"\n",
    "    )\n",
    "    print(f'{processed_prompt[\"pixel_values\"].shape=}')\n",
    "    input_ids = processed_prompt[\"input_ids\"][0].detach().cpu().tolist()\n",
    "\n",
    "    image_outputs = vision(\n",
    "        processed_prompt[\"pixel_values\"].to(\"cuda\"), output_hidden_states=False\n",
    "    )\n",
    "    image_features = image_outputs.last_hidden_state\n",
    "    vision_flat = image_features.view(-1, image_features.size(-1))\n",
    "    precomputed_embeddings = multi_modal_projector(vision_flat)\n",
    "\n",
    "    mm_item = dict(modality=\"IMAGE\", precomputed_embeddings=precomputed_embeddings)\n",
    "    out = llm.generate(input_ids=input_ids, image_data=[mm_item])\n",
    "    print(out[\"text\"])"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "custom_cell_magics": "kql",
   "encoding": "# -*- coding: utf-8 -*-"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
