add vllm

app.py

@@ -8,21 +8,17 @@ import requests
 from huggingface_hub import login
 import torch
 import torch.nn.functional as F
-import spaces
+# import spaces
 import json
-import gradio as gr
 from huggingface_hub import snapshot_download
-import os
-# from loadimg import load_img
-import traceback
+# import traceback
 
 login(os.environ.get("HUGGINGFACE_TOKEN"))
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:80"
 
 repo_id = "mistralai/Pixtral-12B-2409"
-sampling_params = SamplingParams(max_tokens=8192, temperature=0.7)
-max_tokens_per_img = 4096
-max_img_per_msg = 5
-
+max_tokens_per_img = 2048
+max_img_per_msg = 2
 
 title = "# **WIP / DEMO** 🙋🏻‍♂️Welcome to Tonic's Pixtral Model Demo"
 description = """
@@ -40,9 +36,15 @@ with open(f'{model_path}/tekken.json', 'r') as f:
 
 model_name = "mistralai/Pixtral-12B-2409"
 
-sampling_params = SamplingParams(max_tokens=8192)
+llm = LLM(
+    model=model_name,
+    tokenizer_mode="mistral",
+    max_num_batched_tokens=max_img_per_msg * max_tokens_per_img,
+    dtype="float16"
+)
 
-llm = LLM(model=model_name, tokenizer_mode="mistral")
+def clear_cuda_cache():
+    torch.cuda.empty_cache()
 
 def encode_image(image: Image.Image, image_format="PNG") -> str:
     im_file = BytesIO()
@@ -51,11 +53,13 @@ def encode_image(image: Image.Image, image_format="PNG") -> str:
     im_64 = base64.b64encode(im_bytes).decode("utf-8")
     return im_64
 
-def infer(image_url, prompt, progress=gr.Progress(track_tqdm=True)):
+def infer(image_url, prompt, temperature, max_tokens, progress=gr.Progress(track_tqdm=True)):
     if llm is None:
         return "Error: LLM initialization failed. Please try again later."
     
     try:
+        sampling_params = SamplingParams(max_tokens=max_tokens, temperature=temperature)
+        
         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')}"
@@ -68,16 +72,19 @@ def infer(image_url, prompt, progress=gr.Progress(track_tqdm=True)):
         ]
 
         outputs = llm.chat(messages, sampling_params=sampling_params)
-
+        clear_cuda_cache()
         return outputs[0].outputs[0].text
     except Exception as e:
+        clear_cuda_cache()
         return f"Error during inference: {e}"
 
-def compare_images(image1_url, image2_url, prompt, progress=gr.Progress(track_tqdm=True)):
+def compare_images(image1_url, image2_url, prompt, temperature, max_tokens, progress=gr.Progress(track_tqdm=True)):
     if llm is None:
         return "Error: LLM initialization failed. Please try again later."
     
     try:
+        sampling_params = SamplingParams(max_tokens=max_tokens, temperature=temperature)
+        
         image1 = Image.open(BytesIO(requests.get(image1_url).content))
         image2 = Image.open(BytesIO(requests.get(image2_url).content))
         image1 = image1.resize((3844, 2408))
@@ -97,9 +104,10 @@ def compare_images(image1_url, image2_url, prompt, progress=gr.Progress(track_tq
         ]
 
         outputs = llm.chat(messages, sampling_params=sampling_params)
-
+        clear_cuda_cache()
         return outputs[0].outputs[0].text
     except Exception as e:
+        clear_cuda_cache()
         return f"Error during image comparison: {e}"
 
 def calculate_image_similarity(image1_url, image2_url):
@@ -120,9 +128,10 @@ def calculate_image_similarity(image1_url, image2_url):
             embedding2 = llm.model.vision_encoder([image2_tensor])
 
         similarity = F.cosine_similarity(embedding1.mean(dim=0), embedding2.mean(dim=0), dim=0).item()
-
+        clear_cuda_cache()
         return similarity
     except Exception as e:
+        clear_cuda_cache()
         return f"Error during image similarity calculation: {e}"
 
 with gr.Blocks() as demo:
@@ -137,10 +146,12 @@ with gr.Blocks() as demo:
         1. For Image-to-Text Generation:
            - Enter the URL of an image
            - Provide a prompt describing what you want to know about the image
+           - Adjust the temperature and max tokens
            - Click "Generate" to get the model's response
         2. For Image Comparison:
            - Enter URLs for two images you want to compare
            - Provide a prompt asking about the comparison
+           - Adjust the temperature and max tokens
            - Click "Compare" to get the model's analysis
         3. For Image Similarity:
            - Enter URLs for two images you want to compare
@@ -153,20 +164,26 @@ with gr.Blocks() as demo:
             with gr.Row():
                 image_url = gr.Text(label="Image URL")
                 prompt = gr.Text(label="Prompt")
+            with gr.Row():
+                temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, label="Temperature")
+                max_tokens = gr.Number(value=4096, label="Max Tokens")
             generate_button = gr.Button("Generate")
             output = gr.Text(label="Generated Text")
             
-            generate_button.click(infer, inputs=[image_url, prompt], outputs=output)
+            generate_button.click(infer, inputs=[image_url, prompt, temperature, max_tokens], outputs=output)
         
         with gr.TabItem("Image Comparison"):
             with gr.Row():
                 image1_url = gr.Text(label="Image 1 URL")
                 image2_url = gr.Text(label="Image 2 URL")
             comparison_prompt = gr.Text(label="Comparison Prompt")
+            with gr.Row():
+                comparison_temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, label="Temperature")
+                comparison_max_tokens = gr.Number(value=4096, label="Max Tokens")
             compare_button = gr.Button("Compare")
             comparison_output = gr.Text(label="Comparison Result")
             
-            compare_button.click(compare_images, inputs=[image1_url, image2_url, comparison_prompt], outputs=comparison_output)
+            compare_button.click(compare_images, inputs=[image1_url, image2_url, comparison_prompt, comparison_temperature, comparison_max_tokens], outputs=comparison_output)
 
         with gr.TabItem("Image Similarity"):
             with gr.Row():
@@ -187,4 +204,4 @@ with gr.Blocks() as demo:
     gr.Markdown(f"- Patch Size: {params['vision_encoder']['patch_size']}x{params['vision_encoder']['patch_size']}")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()