add vllm

app.py

@@ -41,23 +41,11 @@ 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
-
-@spaces.GPU(duration=120)
-def initialize_llm():
-    global llm
-    if llm is None:
-        try:
-            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},
-                      dtype="float16",
-                      device="cuda" if torch.cuda.is_available() else "cpu")
-        except Exception as e:
-            print(f"Error initializing LLM: {e}")
-            llm = 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})
 
 
 def encode_image(image: Image.Image, image_format="PNG") -> str:
@@ -67,9 +55,8 @@ def encode_image(image: Image.Image, image_format="PNG") -> str:
     im_64 = base64.b64encode(im_bytes).decode("utf-8")
     return im_64
 
-@spaces.GPU(duration=120)
+@spaces.GPU()
 def infer(image_url, prompt, progress=gr.Progress(track_tqdm=True)):
-    initialize_llm()
     if llm is None:
         return "Error: LLM initialization failed. Please try again later."
     
@@ -88,9 +75,8 @@ def infer(image_url, prompt, progress=gr.Progress(track_tqdm=True)):
 
     return outputs[0].outputs[0].text
 
-@spaces.GPU(duration=120)
+@spaces.GPU()
 def compare_images(image1_url, image2_url, prompt, progress=gr.Progress(track_tqdm=True)):
-    initialize_llm()
     if llm is None:
         return "Error: LLM initialization failed. Please try again later."
     
@@ -118,7 +104,6 @@ def compare_images(image1_url, image2_url, prompt, progress=gr.Progress(track_tq
 
 @spaces.GPU(duration=120)
 def calculate_image_similarity(image1_url, image2_url):
-    initialize_llm()
     if llm is None:
         return "Error: LLM initialization failed. Please try again later."
     
@@ -138,7 +123,6 @@ def calculate_image_similarity(image1_url, image2_url):
 
     return similarity
 
-# Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown(title)
     gr.Markdown("## How it works")