add vllm

app.py

@@ -43,15 +43,21 @@ with open(f'{model_path}/tekken.json', 'r') as f:
 
 llm = None
 
-@spaces.GPU()
+@spaces.GPU(duration=120)
 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})
+        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
 
 
 def encode_image(image: Image.Image, image_format="PNG") -> str:
@@ -61,9 +67,12 @@ def encode_image(image: Image.Image, image_format="PNG") -> str:
     im_64 = base64.b64encode(im_bytes).decode("utf-8")
     return im_64
 
-@spaces.GPU()
+@spaces.GPU(duration=120)
 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."
+    
     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')}"
@@ -79,9 +88,12 @@ def infer(image_url, prompt, progress=gr.Progress(track_tqdm=True)):
 
     return outputs[0].outputs[0].text
 
-@spaces.GPU()
+@spaces.GPU(duration=120)
 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."
+    
     image1 = Image.open(BytesIO(requests.get(image1_url).content))
     image2 = Image.open(BytesIO(requests.get(image2_url).content))
     image1 = image1.resize((3844, 2408))
@@ -104,10 +116,12 @@ def compare_images(image1_url, image2_url, prompt, progress=gr.Progress(track_tq
 
     return outputs[0].outputs[0].text
 
-@spaces.GPU()
+@spaces.GPU(duration=120)
 def calculate_image_similarity(image1_url, image2_url):
     initialize_llm()
-    # Load and preprocess images
+    if llm is None:
+        return "Error: LLM initialization failed. Please try again later."
+    
     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