add vllm

app.py

@@ -10,23 +10,11 @@ import torch
 import torch.nn.functional as F
 import spaces
 import json 
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from safetensors import safe_open
 import json
 import gradio as gr
 from PIL import Image
-import numpy as np
 from huggingface_hub import snapshot_download
-# from mistral_common.protocol.instruct.messages import UserMessage, TextChunk, ImageChunk
-# from mistral_common.protocol.instruct.request import ChatCompletionRequest
-# from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
 import spaces
-import math
-from typing import List, Optional, Tuple
-import gc
-# from contextlib import contextmanager
 import os
 from loadimg import load_img
 import traceback
@@ -53,13 +41,17 @@ with open(f'{model_path}/params.json', 'r') as f:
 with open(f'{model_path}/tekken.json', 'r') as f:
     tokenizer_config = json.load(f)
 
+llm = None
+
+def initialize_llm():
+    global llm
+    if llm is None:
+        llm = LLM(model=repo_id,
+                  tokenizer_mode="mistral",
+                  max_model_len=65536,
+                  max_num_batched_tokens=max_img_per_msg * max_tokens_per_img,
+                  limit_mm_per_prompt={"image": max_img_per_msg})
 
-# Initialize the LLM
-llm = LLM(model=repo_id,
-          tokenizer_mode="mistral",
-          max_model_len=65536,
-          max_num_batched_tokens=max_img_per_msg * max_tokens_per_img,
-          limit_mm_per_prompt={"image": max_img_per_msg})
 
 def encode_image(image: Image.Image, image_format="PNG") -> str:
     im_file = BytesIO()
@@ -70,6 +62,7 @@ def encode_image(image: Image.Image, image_format="PNG") -> str:
 
 @spaces.GPU()
 def infer(image_url, prompt, progress=gr.Progress(track_tqdm=True)):
+    initialize_llm()
     image = Image.open(BytesIO(requests.get(image_url).content))
     image = image.resize((3844, 2408))
     new_image_url = f"data:image/png;base64,{encode_image(image, image_format='PNG')}"
@@ -87,6 +80,7 @@ def infer(image_url, prompt, progress=gr.Progress(track_tqdm=True)):
 
 @spaces.GPU()
 def compare_images(image1_url, image2_url, prompt, progress=gr.Progress(track_tqdm=True)):
+    initialize_llm()
     image1 = Image.open(BytesIO(requests.get(image1_url).content))
     image2 = Image.open(BytesIO(requests.get(image2_url).content))
     image1 = image1.resize((3844, 2408))
@@ -111,22 +105,20 @@ def compare_images(image1_url, image2_url, prompt, progress=gr.Progress(track_tq
 
 @spaces.GPU()
 def calculate_image_similarity(image1_url, image2_url):
+    initialize_llm()
     # Load and preprocess images
     image1 = Image.open(BytesIO(requests.get(image1_url).content)).convert('RGB')
     image2 = Image.open(BytesIO(requests.get(image2_url).content)).convert('RGB')
     image1 = image1.resize((224, 224))  # Resize to match model input size
     image2 = image2.resize((224, 224))
 
-    # Convert images to tensors
     image1_tensor = torch.tensor(list(image1.getdata())).view(1, 3, 224, 224).float() / 255.0
     image2_tensor = torch.tensor(list(image2.getdata())).view(1, 3, 224, 224).float() / 255.0
 
-    # Get image embeddings using the vision encoder
     with torch.no_grad():
         embedding1 = llm.model.vision_encoder([image1_tensor])
         embedding2 = llm.model.vision_encoder([image2_tensor])
 
-    # Calculate cosine similarity
     similarity = F.cosine_similarity(embedding1.mean(dim=0), embedding2.mean(dim=0), dim=0).item()
 
     return similarity