add reference code from vllm

app.py

@@ -214,65 +214,68 @@ model = load_model(params, model_path)
 tokenizer = MistralTokenizer.from_model("pixtral")
 
 def preprocess_image(image):
+    if image is None:
+        raise ValueError("No image provided")
     image = image.convert('RGB')
     image = image.resize((params['vision_encoder']['image_size'], params['vision_encoder']['image_size']))
     image_tensor = torch.tensor(np.array(image)).permute(2, 0, 1).unsqueeze(0).float() / 255.0
     return image_tensor
 
-@spaces.GPU
-def generate_text(image, prompt):
-    image_tensor = preprocess_image(image).cuda()
-
-    tokenized = tokenizer.encode_chat_completion(
-        ChatCompletionRequest(
-            messages=[
-                UserMessage(
-                    content=[
-                        TextChunk(text=prompt),
-                        ImageChunk(image=image),
-                    ]
-                )
-            ],
-            model="pixtral",
+@spaces.GPU(duration=120)
+def generate_text(image, prompt, max_tokens):
+    try:
+        image_tensor = preprocess_image(image).cuda()
+
+        tokenized = tokenizer.encode_chat_completion(
+            ChatCompletionRequest(
+                messages=[
+                    UserMessage(
+                        content=[
+                            TextChunk(text=prompt),
+                            ImageChunk(image=image),
+                        ]
+                    )
+                ],
+                model="pixtral",
+            )
         )
-    )
-    input_ids = torch.tensor(tokenized.tokens).unsqueeze(0).cuda()
-
-    # Generate text
-    with torch.no_grad():
-        model.cuda()
-        max_length = 100  # add slider
-        for _ in range(max_length):
-            logits = model(image_tensor, input_ids)
-            next_token_logits = logits[0, -1, :]
-            next_token = torch.argmax(next_token_logits, dim=-1)
-            input_ids = torch.cat([input_ids, next_token.unsqueeze(0).unsqueeze(0)], dim=-1)
-            if next_token.item() == tokenizer.eos_token_id:
-                break
-        model.cpu()
-
-    generated_text = tokenizer.decode(input_ids[0].tolist())
-    return generated_text, len(input_ids[0]), 1  # 1 image processed
-
-@spaces.GPU
+        input_ids = torch.tensor(tokenized.tokens).unsqueeze(0).cuda()
+
+        with torch.no_grad():
+            model.cuda()
+            for _ in range(max_tokens):
+                logits = model(image_tensor, input_ids)
+                next_token_logits = logits[0, -1, :]
+                next_token = torch.argmax(next_token_logits, dim=-1)
+                input_ids = torch.cat([input_ids, next_token.unsqueeze(0).unsqueeze(0)], dim=-1)
+                if next_token.item() == tokenizer.eos_token_id:
+                    break
+            model.cpu()
+
+        generated_text = tokenizer.decode(input_ids[0].tolist())
+        return generated_text, len(input_ids[0]), 1  # 1 image processed
+    except Exception as e:
+        return f"Error: {str(e)}", 0, 0
+
+@spaces.GPU(duration=60)
 def calculate_similarity(image1, image2):
-    # Preprocess images
-    tensor1 = preprocess_image(image1).cuda()
-    tensor2 = preprocess_image(image2).cuda()
+    try:
+        tensor1 = preprocess_image(image1).cuda()
+        tensor2 = preprocess_image(image2).cuda()
 
-    # Generate embeddings
-    with torch.no_grad():
-        model.cuda()
-        embedding1 = model(tensor1).mean(dim=1)  # Average over spatial dimensions
-        embedding2 = model(tensor2).mean(dim=1)
-        model.cpu()
+        with torch.no_grad():
+            model.cuda()
+            embedding1 = model(tensor1).mean(dim=1)  # Average over spatial dimensions
+            embedding2 = model(tensor2).mean(dim=1)
+            model.cpu()
 
-    # Calculate cosine similarity
-    similarity = F.cosine_similarity(embedding1, embedding2).item()
+        similarity = F.cosine_similarity(embedding1, embedding2).item()
 
-    return similarity
+        return similarity
+    except Exception as e:
+        return f"Error: {str(e)}"
 
-with gr.Blocks(theme=gr.themes.Base()) as demo:
+with gr.Blocks() as demo:
     gr.Markdown(title)
     gr.Markdown("## Model Details")
     gr.Markdown(f"- Model Dimension: {params['dim']}")
@@ -287,6 +290,7 @@ with gr.Blocks(theme=gr.themes.Base()) as demo:
     gr.Markdown("1. The image is processed by a Vision Encoder using 2D ROPE (Rotary Position Embedding).")
     gr.Markdown("2. The encoder uses SiLU activation in its feed-forward layers.")
     gr.Markdown("3. The encoded image is used for text generation or similarity comparison.")
+
     gr.Markdown(description)
     
     with gr.Tabs():
@@ -295,6 +299,7 @@ with gr.Blocks(theme=gr.themes.Base()) as demo:
                 with gr.Column():
                     input_image = gr.Image(type="pil", label="Input Image")
                     input_prompt = gr.Textbox(label="Prompt")
+                    max_tokens_slider = gr.Slider(minimum=60, maximum=1600, value=90, step=5, label="Max Tokens")
                     submit_btn = gr.Button("Generate Text")
                 
                 with gr.Column():
@@ -304,7 +309,7 @@ with gr.Blocks(theme=gr.themes.Base()) as demo:
             
             submit_btn.click(
                 fn=generate_text,
-                inputs=[input_image, input_prompt],
+                inputs=[input_image, input_prompt, max_tokens_slider],
                 outputs=[output_text, token_count, image_count]
             )