add vllm

app.py

@@ -53,6 +53,7 @@ def initialize_llm():
         print("LLM initialization failed:", e)
         return None
 
+sampling_params = SamplingParams(max_tokens=8192)
 llm = initialize_llm()
 
 def encode_image(image: Image.Image, image_format="PNG") -> str:
@@ -67,68 +68,77 @@ def infer(image_url, prompt, progress=gr.Progress(track_tqdm=True)):
     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')}"
+    try:
+        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')}"
 
-    messages = [
-        {
-            "role": "user",
-            "content": [{"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": new_image_url}}]
-        },
-    ]
+        messages = [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": new_image_url}}]
+            },
+        ]
 
-    outputs = llm.chat(messages, sampling_params=sampling_params)
+        outputs = llm.chat(messages, sampling_params=sampling_params)
 
-    return outputs[0].outputs[0].text
+        return outputs[0].outputs[0].text
+    except Exception as e:
+        return f"Error during inference: {e}"
 
 @spaces.GPU()
 def compare_images(image1_url, image2_url, prompt, progress=gr.Progress(track_tqdm=True)):
     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))
-    image2 = image2.resize((3844, 2408))
-    new_image1_url = f"data:image/png;base64,{encode_image(image1, image_format='PNG')}"
-    new_image2_url = f"data:image/png;base64,{encode_image(image2, image_format='PNG')}"
-
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "text", "text": prompt},
-                {"type": "image_url", "image_url": {"url": new_image1_url}},
-                {"type": "image_url", "image_url": {"url": new_image2_url}}
-            ]
-        },
-    ]
-
-    outputs = llm.chat(messages, sampling_params=sampling_params)
-
-    return outputs[0].outputs[0].text
+    try:
+        image1 = Image.open(BytesIO(requests.get(image1_url).content))
+        image2 = Image.open(BytesIO(requests.get(image2_url).content))
+        image1 = image1.resize((3844, 2408))
+        image2 = image2.resize((3844, 2408))
+        new_image1_url = f"data:image/png;base64,{encode_image(image1, image_format='PNG')}"
+        new_image2_url = f"data:image/png;base64,{encode_image(image2, image_format='PNG')}"
+
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt},
+                    {"type": "image_url", "image_url": {"url": new_image1_url}},
+                    {"type": "image_url", "image_url": {"url": new_image2_url}}
+                ]
+            },
+        ]
+
+        outputs = llm.chat(messages, sampling_params=sampling_params)
+
+        return outputs[0].outputs[0].text
+    except Exception as e:
+        return f"Error during image comparison: {e}"
 
 @spaces.GPU()
 def calculate_image_similarity(image1_url, image2_url):
     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
-    image2 = image2.resize((224, 224))
+    try:
+        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))
 
-    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
+        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
 
-    with torch.no_grad():
-        embedding1 = llm.model.vision_encoder([image1_tensor])
-        embedding2 = llm.model.vision_encoder([image2_tensor])
+        with torch.no_grad():
+            embedding1 = llm.model.vision_encoder([image1_tensor])
+            embedding2 = llm.model.vision_encoder([image2_tensor])
 
-    similarity = F.cosine_similarity(embedding1.mean(dim=0), embedding2.mean(dim=0), dim=0).item()
+        similarity = F.cosine_similarity(embedding1.mean(dim=0), embedding2.mean(dim=0), dim=0).item()
 
-    return similarity
+        return similarity
+    except Exception as e:
+        return f"Error during image similarity calculation: {e}"
 
 with gr.Blocks() as demo:
     gr.Markdown(title)