Spaces:
Running
Running
File size: 7,553 Bytes
72cd7d2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 |
import gradio as gr
import torch
import numpy as np
from PIL import Image
import os
import json
import base64
from io import BytesIO
import requests
from typing import Dict, List, Any, Optional
from transformers.pipelines import pipeline
# MCP imports
from modelcontextprotocol.server import Server
from modelcontextprotocol.server.gradio import GradioServerTransport
from modelcontextprotocol.types import (
CallToolRequestSchema,
ErrorCode,
ListToolsRequestSchema,
McpError,
)
# Initialize the model
model = pipeline("image-feature-extraction", model="nomic-ai/nomic-embed-vision-v1.5", trust_remote_code=True)
# Function to generate embeddings from an image
def generate_embedding(image):
if image is None:
return None
# Convert to PIL Image if needed
if not isinstance(image, Image.Image):
image = Image.fromarray(image)
try:
# Generate embedding using the transformers pipeline
result = model(image)
# Process the result based on its type
embedding_list = None
# Handle different possible output types
if isinstance(result, torch.Tensor):
embedding_list = result.detach().cpu().numpy().flatten().tolist()
elif isinstance(result, np.ndarray):
embedding_list = result.flatten().tolist()
elif isinstance(result, list):
# If it's a list of tensors or arrays
if result and isinstance(result[0], (torch.Tensor, np.ndarray)):
embedding_list = result[0].flatten().tolist() if hasattr(result[0], 'flatten') else result[0]
else:
embedding_list = result
else:
# Try to convert to a list as a last resort
try:
if result is not None:
embedding_list = list(result)
else:
print("Result is None")
return None
except:
print(f"Couldn't convert result of type {type(result)} to list")
return None
# Ensure we have a valid embedding list
if embedding_list is None:
return None
# Calculate embedding dimension
embedding_dim = len(embedding_list)
return {
"embedding": embedding_list,
"dimension": embedding_dim
}
except Exception as e:
print(f"Error generating embedding: {str(e)}")
return None
# Gradio Interface
with gr.Blocks() as demo:
gr.Markdown("# Nomic Vision Embedding Model (nomic-ai/nomic-embed-vision-v1.5)")
gr.Markdown("Upload an image to generate embeddings using the Nomic Vision model.")
with gr.Row():
with gr.Column():
input_image = gr.Image(type="pil", label="Input Image")
embed_btn = gr.Button("Generate Embedding")
with gr.Column():
embedding_json = gr.JSON(label="Embedding Output")
embedding_dim = gr.Textbox(label="Embedding Dimension")
def update_embedding(img):
result = generate_embedding(img)
if result is None:
return {
embedding_json: None,
embedding_dim: "No embedding generated"
}
return {
embedding_json: result,
embedding_dim: f"Dimension: {len(result['embedding'])}"
}
embed_btn.click(
fn=update_embedding,
inputs=[input_image],
outputs=[embedding_json, embedding_dim]
)
# MCP Server Implementation
class NomicEmbeddingServer:
def __init__(self):
self.server = Server(
{
"name": "nomic-embedding-server",
"version": "0.1.0",
},
{
"capabilities": {
"tools": {},
},
}
)
self.setup_tool_handlers()
# Error handling
self.server.onerror = lambda error: print(f"[MCP Error] {error}")
def setup_tool_handlers(self):
self.server.set_request_handler(ListToolsRequestSchema, self.handle_list_tools)
self.server.set_request_handler(CallToolRequestSchema, self.handle_call_tool)
async def handle_list_tools(self, request):
return {
"tools": [
{
"name": "embed_image",
"description": "Generate embeddings for an image using nomic-ai/nomic-embed-vision-v1.5",
"inputSchema": {
"type": "object",
"properties": {
"image_url": {
"type": "string",
"description": "URL of the image to embed",
},
"image_data": {
"type": "string",
"description": "Base64-encoded image data (alternative to image_url)",
},
},
"anyOf": [
{"required": ["image_url"]},
{"required": ["image_data"]},
],
},
}
]
}
async def handle_call_tool(self, request):
if request.params.name != "embed_image":
raise McpError(
ErrorCode.MethodNotFound,
f"Unknown tool: {request.params.name}"
)
args = request.params.arguments
try:
# Handle image from URL
if "image_url" in args:
response = requests.get(args["image_url"])
image = Image.open(BytesIO(response.content))
# Handle image from base64 data
elif "image_data" in args:
image_data = base64.b64decode(args["image_data"])
image = Image.open(BytesIO(image_data))
else:
raise McpError(
ErrorCode.InvalidParams,
"Either image_url or image_data must be provided"
)
# Generate embedding
result = generate_embedding(image)
return {
"content": [
{
"type": "text",
"text": json.dumps(result, indent=2),
}
]
}
except Exception as e:
return {
"content": [
{
"type": "text",
"text": f"Error generating embedding: {str(e)}",
}
],
"isError": True,
}
# Initialize and run the MCP server
embedding_server = NomicEmbeddingServer()
# Connect the MCP server to the Gradio app
transport = GradioServerTransport(demo)
embedding_server.server.connect(transport)
# Launch the Gradio app
if __name__ == "__main__":
# For Huggingface Spaces, we need to specify the server name and port
demo.launch(server_name="0.0.0.0", server_port=7860) |