From 635801d51580a8f89c377494191b5b56346d841b Mon Sep 17 00:00:00 2001 From: rltakashige Date: Mon, 30 Mar 2026 11:52:19 +0100 Subject: [PATCH] Add multimodality! (#1802) ## Motivation Images! TODO (in a future PR): Add audio and video support. ## Test Plan ### Manual Testing image image image (And batching also works) --------- Co-authored-by: David Hind --- .mlx_typings/mlx_vlm/__init__.pyi | 0 .mlx_typings/mlx_vlm/prompt_utils.pyi | 12 + .mlx_typings/mlx_vlm/utils.pyi | 15 + .mlx_typings/safetensors/__init__.pyi | 8 + dashboard/src/lib/stores/app.svelte.ts | 44 +- flake.lock | 18 +- pyproject.toml | 5 +- .../mlx-community--Kimi-K2.5.toml | 8 +- ...-community--Qwen3-VL-4B-Instruct-4bit.toml | 12 + src/exo/api/adapters/chat_completions.py | 64 +- src/exo/api/adapters/claude.py | 37 +- src/exo/api/adapters/ollama.py | 27 +- src/exo/api/adapters/responses.py | 35 +- src/exo/api/main.py | 97 ++- src/exo/api/types/__init__.py | 2 + src/exo/api/types/api.py | 10 +- src/exo/api/types/claude_api.py | 2 +- src/exo/api/types/ollama_api.py | 1 + src/exo/api/types/openai_responses.py | 10 +- src/exo/download/impl_shard_downloader.py | 45 +- src/exo/shared/models/model_cards.py | 93 ++- src/exo/shared/types/chunks.py | 1 + src/exo/shared/types/text_generation.py | 6 +- src/exo/worker/engines/mlx/cache.py | 54 ++ .../engines/mlx/generator/batch_generate.py | 68 +- .../worker/engines/mlx/generator/generate.py | 94 ++- src/exo/worker/engines/mlx/utils_mlx.py | 20 +- src/exo/worker/engines/mlx/vision.py | 589 ++++++++++++++++++ src/exo/worker/main.py | 57 +- src/exo/worker/plan.py | 17 +- .../runner/llm_inference/batch_generator.py | 5 + src/exo/worker/runner/llm_inference/runner.py | 26 +- .../unittests/test_mlx/test_tokenizers.py | 4 +- .../test_runner/test_event_ordering.py | 4 +- tests/test_vision_cache.py | 63 ++ uv.lock | 429 ++++++++----- 36 files changed, 1714 insertions(+), 268 deletions(-) create mode 100644 .mlx_typings/mlx_vlm/__init__.pyi create mode 100644 .mlx_typings/mlx_vlm/prompt_utils.pyi create mode 100644 .mlx_typings/mlx_vlm/utils.pyi create mode 100644 .mlx_typings/safetensors/__init__.pyi create mode 100644 resources/inference_model_cards/mlx-community--Qwen3-VL-4B-Instruct-4bit.toml create mode 100644 src/exo/worker/engines/mlx/vision.py create mode 100644 tests/test_vision_cache.py diff --git a/.mlx_typings/mlx_vlm/__init__.pyi b/.mlx_typings/mlx_vlm/__init__.pyi new file mode 100644 index 00000000..e69de29b diff --git a/.mlx_typings/mlx_vlm/prompt_utils.pyi b/.mlx_typings/mlx_vlm/prompt_utils.pyi new file mode 100644 index 00000000..d21d0b9a --- /dev/null +++ b/.mlx_typings/mlx_vlm/prompt_utils.pyi @@ -0,0 +1,12 @@ +from typing import Any + +def get_message_json( + model_name: str, + prompt: str, + role: str = "user", + skip_image_token: bool = False, + skip_audio_token: bool = False, + num_images: int = 0, + num_audios: int = 0, + **kwargs: Any, +) -> dict[str, Any]: ... diff --git a/.mlx_typings/mlx_vlm/utils.pyi b/.mlx_typings/mlx_vlm/utils.pyi new file mode 100644 index 00000000..03523728 --- /dev/null +++ b/.mlx_typings/mlx_vlm/utils.pyi @@ -0,0 +1,15 @@ +from pathlib import Path +from typing import Any + +class ImageProcessor: + def preprocess( + self, images: list[dict[str, Any]], **kwargs: Any + ) -> dict[str, Any]: ... + def __call__(self, **kwargs: Any) -> dict[str, Any]: ... + +def load_image_processor( + model_path: str | Path, **kwargs: Any +) -> ImageProcessor | None: ... +def load_processor( + model_path: str | Path, add_detokenizer: bool = ..., **kwargs: Any +) -> ImageProcessor: ... diff --git a/.mlx_typings/safetensors/__init__.pyi b/.mlx_typings/safetensors/__init__.pyi new file mode 100644 index 00000000..a45a06e6 --- /dev/null +++ b/.mlx_typings/safetensors/__init__.pyi @@ -0,0 +1,8 @@ +from typing import Any, Self + +class safe_open: + def __init__(self, filename: str, framework: str = "pt") -> None: ... + def __enter__(self) -> Self: ... + def __exit__(self, *args: Any) -> None: ... + def keys(self) -> list[str]: ... + def get_tensor(self, name: str) -> Any: ... diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts index 65c7e20c..b43776cc 100644 --- a/dashboard/src/lib/stores/app.svelte.ts +++ b/dashboard/src/lib/stores/app.svelte.ts @@ -2344,7 +2344,49 @@ class AppStore { const apiMessages = [ systemPrompt, ...targetConversation.messages.slice(0, -1).map((m) => { - // Build content including any text file attachments + // Check if this message has image attachments + const imageAttachments = m.attachments?.filter( + (a) => a.type === "image" && a.preview, + ); + + if (imageAttachments && imageAttachments.length > 0) { + // Build multimodal content array (OpenAI vision format) + const contentParts: Array< + | { type: "text"; text: string } + | { type: "image_url"; image_url: { url: string } } + > = []; + + // Add image parts first + for (const img of imageAttachments) { + if (img.preview) { + contentParts.push({ + type: "image_url", + image_url: { url: img.preview }, + }); + } + } + + // Build text content including any text file attachments + let textContent = m.content; + if (m.attachments) { + for (const attachment of m.attachments) { + if (attachment.type === "text" && attachment.content) { + textContent += `\n\n[File: ${attachment.name}]\n\`\`\`\n${attachment.content}\n\`\`\``; + } + } + } + + if (textContent) { + contentParts.push({ type: "text", text: textContent }); + } + + return { + role: m.role, + content: contentParts, + }; + } + + // Text-only message (original path) let msgContent = m.content; // Add text attachments as context diff --git a/flake.lock b/flake.lock index 86fb6e60..6c589f9e 100644 --- a/flake.lock +++ b/flake.lock @@ -164,11 +164,11 @@ ] }, "locked": { - "lastModified": 1763662255, - "narHash": "sha256-4bocaOyLa3AfiS8KrWjZQYu+IAta05u3gYZzZ6zXbT0=", + "lastModified": 1773870109, + "narHash": "sha256-ZoTdqZP03DcdoyxvpFHCAek4bkPUTUPUF3oCCgc3dP4=", "owner": "pyproject-nix", "repo": "build-system-pkgs", - "rev": "042904167604c681a090c07eb6967b4dd4dae88c", + "rev": "b6e74f433b02fa4b8a7965ee24680f4867e2926f", "type": "github" }, "original": { @@ -184,11 +184,11 @@ ] }, "locked": { - "lastModified": 1764134915, - "narHash": "sha256-xaKvtPx6YAnA3HQVp5LwyYG1MaN4LLehpQI8xEdBvBY=", + "lastModified": 1774498001, + "narHash": "sha256-wTfdyzzrmpuqt4TQQNqilF91v0m5Mh1stNy9h7a/WK4=", "owner": "pyproject-nix", "repo": "pyproject.nix", - "rev": "2c8df1383b32e5443c921f61224b198a2282a657", + "rev": "794afa6eb588b498344f2eaa36ab1ceb7e6b0b09", "type": "github" }, "original": { @@ -280,11 +280,11 @@ ] }, "locked": { - "lastModified": 1767701098, - "narHash": "sha256-CJhKZnWb3gumR9oTRjFvCg/6lYTGbZRU7xtvcyWIRwU=", + "lastModified": 1774490495, + "narHash": "sha256-a9WmQWj8fF7BctZGCoyzpUjP6GJw8H+lxl+zxpGnETk=", "owner": "pyproject-nix", "repo": "uv2nix", - "rev": "9d357f0d2ce6f5f35ec7959d7e704452352eb4da", + "rev": "18ae62fc5e389e3069854a7c66455c22e31708fc", "type": "github" }, "original": { diff --git a/pyproject.toml b/pyproject.toml index 81beb828..9895e0c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,7 @@ dependencies = [ "fastapi>=0.116.1", "filelock>=3.18.0", "rustworkx>=0.17.1", - "huggingface-hub>=0.33.4", + "huggingface-hub>=1.8.0", "psutil>=7.0.0", "loguru>=0.7.3", "exo_pyo3_bindings", # rust bindings @@ -29,6 +29,8 @@ dependencies = [ "python-multipart>=0.0.21", "msgspec>=0.19.0", "zstandard>=0.23.0", + "mlx-vlm>=0.3.11", + "transformers>=5.0.0,<5.4.0", ] [project.scripts] @@ -114,6 +116,7 @@ root = "src" required-version = ">=0.8.6" prerelease = "allow" environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"] +extra-build-dependencies = { "miniaudio" = ["setuptools", "cffi", "pycparser"] } ### # ruff configuration diff --git a/resources/inference_model_cards/mlx-community--Kimi-K2.5.toml b/resources/inference_model_cards/mlx-community--Kimi-K2.5.toml index 7bb689cc..a2eea6c3 100644 --- a/resources/inference_model_cards/mlx-community--Kimi-K2.5.toml +++ b/resources/inference_model_cards/mlx-community--Kimi-K2.5.toml @@ -7,7 +7,13 @@ tasks = ["TextGeneration"] family = "kimi" quantization = "" base_model = "Kimi K2.5" -capabilities = ["text", "thinking", "thinking_toggle"] +capabilities = ["text", "thinking", "thinking_toggle", "vision"] [storage_size] in_bytes = 662498705408 + +[vision] +image_token_id = 163605 +model_type = "kimi_vl" +weights_repo = "davehind/Kimi-K2.5-vision" +processor_repo = "moonshotai/Kimi-K2.5" diff --git a/resources/inference_model_cards/mlx-community--Qwen3-VL-4B-Instruct-4bit.toml b/resources/inference_model_cards/mlx-community--Qwen3-VL-4B-Instruct-4bit.toml new file mode 100644 index 00000000..b1cd18f6 --- /dev/null +++ b/resources/inference_model_cards/mlx-community--Qwen3-VL-4B-Instruct-4bit.toml @@ -0,0 +1,12 @@ +model_id = "mlx-community/Qwen3-VL-4B-Instruct-4bit" +n_layers = 36 +hidden_size = 2560 +supports_tensor = true +tasks = ["TextGeneration"] +family = "qwen" +quantization = "4bit" +base_model = "Qwen3-VL 4B" +capabilities = ["text", "thinking", "vision"] + +[storage_size] +in_bytes = 3340000000 diff --git a/src/exo/api/adapters/chat_completions.py b/src/exo/api/adapters/chat_completions.py index 38939151..17ebbf6b 100644 --- a/src/exo/api/adapters/chat_completions.py +++ b/src/exo/api/adapters/chat_completions.py @@ -1,5 +1,7 @@ """OpenAI Chat Completions API adapter for converting requests/responses.""" +import base64 +import re import time from collections.abc import AsyncGenerator from typing import Any @@ -7,6 +9,7 @@ from typing import Any from exo.api.types import ( ChatCompletionChoice, ChatCompletionMessage, + ChatCompletionMessageImageUrl, ChatCompletionMessageText, ChatCompletionRequest, ChatCompletionResponse, @@ -19,6 +22,7 @@ from exo.api.types import ( ToolCall, Usage, ) +from exo.download.download_utils import create_http_session from exo.shared.types.chunks import ( ErrorChunk, PrefillProgressChunk, @@ -33,25 +37,65 @@ from exo.shared.types.text_generation import ( ) -def chat_request_to_text_generation( +def extract_base64_from_data_url(data_url: str) -> str: + match = re.match(r"data:[^;]+;base64,(.+)", data_url) + if match: + return match.group(1) + return data_url + + +async def fetch_image_url(url: str) -> str: + headers = {"User-Agent": "exo/1.0"} + async with ( + create_http_session(timeout_profile="short") as session, + session.get(url, headers=headers) as resp, + ): + resp.raise_for_status() + data = await resp.read() + return base64.b64encode(data).decode("ascii") + + +async def chat_request_to_text_generation( request: ChatCompletionRequest, ) -> TextGenerationTaskParams: instructions: str | None = None input_messages: list[InputMessage] = [] chat_template_messages: list[dict[str, Any]] = [] + images: list[str] = [] for msg in request.messages: # Normalize content to string content: str + has_images = False if msg.content is None: content = "" elif isinstance(msg.content, str): content = msg.content elif isinstance(msg.content, ChatCompletionMessageText): content = msg.content.text + elif isinstance(msg.content, ChatCompletionMessageImageUrl): + url = msg.content.image_url.get("url", "") + if url: + if url.startswith(("http://", "https://")): + images.append(await fetch_image_url(url)) + else: + images.append(extract_base64_from_data_url(url)) + has_images = True + content = "" else: - # List of ChatCompletionMessageText - content = "\n".join(item.text for item in msg.content) + text_parts: list[str] = [] + for part in msg.content: + if isinstance(part, ChatCompletionMessageText): + text_parts.append(part.text) + else: + url = part.image_url.get("url", "") + if url: + if url.startswith(("http://", "https://")): + images.append(await fetch_image_url(url)) + else: + images.append(extract_base64_from_data_url(url)) + has_images = True + content = "\n".join(text_parts) # Extract system message as instructions if msg.role == "system": @@ -75,7 +119,20 @@ def chat_request_to_text_generation( # Build full message dict for chat template (preserves tool_calls etc.) # Normalize content for model_dump + if has_images: + multimodal_content: list[dict[str, Any]] = [] + assert isinstance(msg.content, list) + for part in msg.content: + if isinstance(part, ChatCompletionMessageText): + multimodal_content.append({"type": "text", "text": part.text}) + else: + multimodal_content.append({"type": "image"}) + chat_template_messages.append( + {"role": msg.role, "content": multimodal_content} + ) + continue msg_copy = msg.model_copy(update={"content": content}) + dumped: dict[str, Any] = msg_copy.model_dump(exclude_none=True) chat_template_messages.append(dumped) @@ -107,6 +164,7 @@ def chat_request_to_text_generation( min_p=request.min_p, repetition_penalty=request.repetition_penalty, repetition_context_size=request.repetition_context_size, + images=images, ) diff --git a/src/exo/api/adapters/claude.py b/src/exo/api/adapters/claude.py index 7ce94200..0b53e330 100644 --- a/src/exo/api/adapters/claude.py +++ b/src/exo/api/adapters/claude.py @@ -11,6 +11,7 @@ from exo.api.types.claude_api import ( ClaudeContentBlockDeltaEvent, ClaudeContentBlockStartEvent, ClaudeContentBlockStopEvent, + ClaudeImageBlock, ClaudeInputJsonDelta, ClaudeMessageDelta, ClaudeMessageDeltaEvent, @@ -61,7 +62,9 @@ def _extract_tool_result_text(block: ClaudeToolResultBlock) -> str: return "" if isinstance(block.content, str): return block.content - return "".join(sub_block.text for sub_block in block.content) + return "".join( + sub.text for sub in block.content if isinstance(sub, ClaudeTextBlock) + ) # Matches "x-anthropic-billing-header: ...;" (with optional trailing newline) @@ -86,6 +89,7 @@ def claude_request_to_text_generation( # Handle system message instructions: str | None = None chat_template_messages: list[dict[str, Any]] = [] + images: list[str] = [] if request.system: if isinstance(request.system, str): @@ -109,10 +113,18 @@ def claude_request_to_text_generation( thinking_parts: list[str] = [] tool_calls: list[dict[str, Any]] = [] tool_results: list[ClaudeToolResultBlock] = [] + has_images = False for block in msg.content: if isinstance(block, ClaudeTextBlock): text_parts.append(block.text) + elif isinstance(block, ClaudeImageBlock): + if block.source.type == "base64" and block.source.data: + images.append(block.source.data) + has_images = True + elif block.source.type == "url" and block.source.url: + images.append(block.source.url) + has_images = True elif isinstance(block, ClaudeThinkingBlock): thinking_parts.append(block.thinking) elif isinstance(block, ClaudeToolUseBlock): @@ -126,8 +138,17 @@ def claude_request_to_text_generation( }, } ) - elif isinstance(block, ClaudeToolResultBlock): + else: tool_results.append(block) + if isinstance(block.content, list): + for sub in block.content: + if isinstance(sub, ClaudeImageBlock): + if sub.source.type == "base64" and sub.source.data: + images.append(sub.source.data) + has_images = True + elif sub.source.type == "url" and sub.source.url: + images.append(sub.source.url) + has_images = True content = "".join(text_parts) reasoning_content = "".join(thinking_parts) if thinking_parts else None @@ -155,6 +176,17 @@ def claude_request_to_text_generation( "content": _extract_tool_result_text(tr), } ) + elif has_images: + multimodal_content: list[dict[str, Any]] = [] + for block in msg.content: + if isinstance(block, ClaudeTextBlock): + multimodal_content.append({"type": "text", "text": block.text}) + elif isinstance(block, ClaudeImageBlock): + multimodal_content.append({"type": "image"}) + chat_msg = {"role": msg.role, "content": multimodal_content} + if reasoning_content: + chat_msg["reasoning_content"] = reasoning_content + chat_template_messages.append(chat_msg) else: chat_msg = {"role": msg.role, "content": content} if reasoning_content: @@ -197,6 +229,7 @@ def claude_request_to_text_generation( chat_template_messages=chat_template_messages if chat_template_messages else None, + images=images, ) diff --git a/src/exo/api/adapters/ollama.py b/src/exo/api/adapters/ollama.py index 3eeeb46c..76436370 100644 --- a/src/exo/api/adapters/ollama.py +++ b/src/exo/api/adapters/ollama.py @@ -82,10 +82,15 @@ def ollama_request_to_text_generation( instructions: str | None = None input_messages: list[InputMessage] = [] chat_template_messages: list[dict[str, Any]] = [] + images: list[str] = [] tool_message_index = 0 for msg in request.messages: content = msg.content or "" + has_images = False + if msg.images: + images.extend(msg.images) + has_images = True if msg.role == "system": if instructions is None: @@ -100,6 +105,16 @@ def ollama_request_to_text_generation( ): input_messages.append(InputMessage(role=msg.role, content=content)) + if has_images: + multimodal: list[dict[str, Any]] = [ + {"type": "image"} for _ in (msg.images or []) + ] + if content: + multimodal.append({"type": "text", "text": content}) + chat_template_messages.append({"role": msg.role, "content": multimodal}) + if msg.role in ("user", "assistant"): + input_messages.append(InputMessage(role=msg.role, content=content)) + continue dumped: dict[str, Any] = {"role": msg.role, "content": content} if msg.thinking is not None: dumped["thinking"] = msg.thinking @@ -152,6 +167,7 @@ def ollama_request_to_text_generation( chat_template_messages=chat_template_messages if chat_template_messages else None, + images=images, ) @@ -311,9 +327,17 @@ def ollama_generate_request_to_text_generation( ) -> TextGenerationTaskParams: """Convert Ollama generate request to exo's internal text generation format.""" chat_template_messages: list[dict[str, Any]] = [] + images: list[str] = [] if request.system: chat_template_messages.append({"role": "system", "content": request.system}) - chat_template_messages.append({"role": "user", "content": request.prompt}) + if request.images: + images.extend(request.images) + multimodal: list[dict[str, Any]] = [{"type": "image"} for _ in request.images] + if request.prompt: + multimodal.append({"type": "text", "text": request.prompt}) + chat_template_messages.append({"role": "user", "content": multimodal}) + else: + chat_template_messages.append({"role": "user", "content": request.prompt}) options = request.options return TextGenerationTaskParams( @@ -331,6 +355,7 @@ def ollama_generate_request_to_text_generation( chat_template_messages=chat_template_messages if chat_template_messages else None, + images=images, ) diff --git a/src/exo/api/adapters/responses.py b/src/exo/api/adapters/responses.py index 014aeaf0..37b72558 100644 --- a/src/exo/api/adapters/responses.py +++ b/src/exo/api/adapters/responses.py @@ -4,6 +4,10 @@ from collections.abc import AsyncGenerator from itertools import count from typing import Any +from exo.api.adapters.chat_completions import ( + extract_base64_from_data_url, + fetch_image_url, +) from exo.api.types import Usage from exo.api.types.openai_responses import ( FunctionCallInputItem, @@ -16,6 +20,7 @@ from exo.api.types.openai_responses import ( ResponseFunctionCallArgumentsDoneEvent, ResponseFunctionCallItem, ResponseInProgressEvent, + ResponseInputImagePart, ResponseInputMessage, ResponseItem, ResponseMessageItem, @@ -58,19 +63,23 @@ def _extract_content(content: str | list[ResponseContentPart]) -> str: """Extract plain text from a content field that may be a string or list of parts.""" if isinstance(content, str): return content - return "".join(part.text for part in content) + return "".join( + part.text for part in content if not isinstance(part, ResponseInputImagePart) + ) -def responses_request_to_text_generation( +async def responses_request_to_text_generation( request: ResponsesRequest, ) -> TextGenerationTaskParams: input_value: list[InputMessage] built_chat_template: list[dict[str, Any]] | None = None + images: list[str] = [] if isinstance(request.input, str): input_value = [InputMessage(role="user", content=request.input)] else: input_messages: list[InputMessage] = [] chat_template_messages: list[dict[str, Any]] = [] + has_images = False if request.instructions is not None: chat_template_messages.append( @@ -80,12 +89,33 @@ def responses_request_to_text_generation( for item in request.input: if isinstance(item, ResponseInputMessage): content = _extract_content(item.content) + if isinstance(item.content, list): + for part in item.content: + if isinstance(part, ResponseInputImagePart) and part.image_url: + url = part.image_url + if url.startswith(("http://", "https://")): + images.append(await fetch_image_url(url)) + else: + images.append(extract_base64_from_data_url(url)) + has_images = True if item.role in ("user", "assistant", "developer"): input_messages.append(InputMessage(role=item.role, content=content)) if item.role == "system": chat_template_messages.append( {"role": "system", "content": content} ) + elif has_images: + multimodal: list[dict[str, Any]] = [] + if isinstance(item.content, list): + for part in item.content: + if isinstance(part, ResponseInputImagePart): + multimodal.append({"type": "image"}) + elif hasattr(part, "text"): + multimodal.append({"type": "text", "text": part.text}) + chat_template_messages.append( + {"role": item.role, "content": multimodal} + ) + has_images = False else: chat_template_messages.append( {"role": item.role, "content": content} @@ -165,6 +195,7 @@ def responses_request_to_text_generation( chat_template_messages=built_chat_template or request.chat_template_messages, reasoning_effort=resolved_effort, enable_thinking=resolved_thinking, + images=images, ) diff --git a/src/exo/api/main.py b/src/exo/api/main.py index b09e8504..cc3a159b 100644 --- a/src/exo/api/main.py +++ b/src/exo/api/main.py @@ -1,5 +1,6 @@ import base64 import contextlib +import hashlib import json import random import time @@ -24,6 +25,7 @@ from loguru import logger from exo.api.adapters.chat_completions import ( chat_request_to_text_generation, collect_chat_response, + fetch_image_url, generate_chat_stream, ) from exo.api.adapters.claude import ( @@ -170,6 +172,7 @@ from exo.shared.types.events import ( ) from exo.shared.types.memory import Memory from exo.shared.types.state import State +from exo.shared.types.text_generation import TextGenerationTaskParams from exo.shared.types.worker.downloads import DownloadCompleted from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta from exo.shared.types.worker.shards import Sharding @@ -720,18 +723,79 @@ class API: "TODO: we should send a notification to the user to download the model" ) + _sent_image_hashes: set[str] = set() + + async def _send_text_generation_with_images( + self, task_params: TextGenerationTaskParams + ) -> TextGeneration: + images = task_params.images + if not images: + command = TextGeneration(task_params=task_params) + await self._send(command) + return command + + hashes = [hashlib.sha256(img.encode("ascii")).hexdigest() for img in images] + + cached_hashes: dict[int, str] = {} + new_images: list[tuple[int, str]] = [] + for idx, (img, h) in enumerate(zip(images, hashes, strict=True)): + if h in self._sent_image_hashes: + cached_hashes[idx] = h + else: + self._sent_image_hashes.add(h) + new_images.append((idx, img)) + + if not new_images: + task_params = task_params.model_copy( + update={"images": [], "image_hashes": cached_hashes} + ) + command = TextGeneration(task_params=task_params) + await self._send(command) + return command + + all_chunks: list[tuple[int, str]] = [] + for img_idx, img_data in new_images: + for i in range(0, len(img_data), EXO_MAX_CHUNK_SIZE): + all_chunks.append((img_idx, img_data[i : i + EXO_MAX_CHUNK_SIZE])) + + task_params = task_params.model_copy( + update={ + "images": [], + "image_hashes": cached_hashes, + "total_input_chunks": len(all_chunks), + "image_count": len(new_images), + } + ) + command = TextGeneration(task_params=task_params) + + for global_idx, (img_idx, chunk_data) in enumerate(all_chunks): + await self._send( + SendInputChunk( + chunk=InputImageChunk( + model=task_params.model, + command_id=command.command_id, + data=chunk_data, + chunk_index=global_idx, + total_chunks=len(all_chunks), + image_index=img_idx, + ) + ) + ) + + await self._send(command) + return command + async def chat_completions( self, payload: ChatCompletionRequest ) -> ChatCompletionResponse | StreamingResponse: """OpenAI Chat Completions API - adapter.""" - task_params = chat_request_to_text_generation(payload) + task_params = await chat_request_to_text_generation(payload) resolved_model = await self._resolve_and_validate_text_model( ModelId(task_params.model) ) task_params = task_params.model_copy(update={"model": resolved_model}) - command = TextGeneration(task_params=task_params) - await self._send(command) + command = await self._send_text_generation_with_images(task_params) if payload.stream: return StreamingResponse( @@ -758,7 +822,7 @@ class API: async def bench_chat_completions( self, payload: BenchChatCompletionRequest ) -> BenchChatCompletionResponse: - task_params = chat_request_to_text_generation(payload) + task_params = await chat_request_to_text_generation(payload) resolved_model = await self._resolve_and_validate_text_model( ModelId(task_params.model) ) @@ -766,8 +830,7 @@ class API: task_params = task_params.model_copy(update={"stream": False, "bench": True}) - command = TextGeneration(task_params=task_params) - await self._send(command) + command = await self._send_text_generation_with_images(task_params) return await self._collect_text_generation_with_stats(command.command_id) @@ -1324,13 +1387,20 @@ class API: ) -> ClaudeMessagesResponse | StreamingResponse: """Claude Messages API - adapter.""" task_params = claude_request_to_text_generation(payload) + if task_params.images: + resolved_images: list[str] = [] + for img in task_params.images: + if img.startswith(("http://", "https://")): + resolved_images.append(await fetch_image_url(img)) + else: + resolved_images.append(img) + task_params = task_params.model_copy(update={"images": resolved_images}) resolved_model = await self._resolve_and_validate_text_model( ModelId(task_params.model) ) task_params = task_params.model_copy(update={"model": resolved_model}) - command = TextGeneration(task_params=task_params) - await self._send(command) + command = await self._send_text_generation_with_images(task_params) if payload.stream: return StreamingResponse( @@ -1360,12 +1430,11 @@ class API: self, payload: ResponsesRequest ) -> ResponsesResponse | StreamingResponse: """OpenAI Responses API.""" - task_params = responses_request_to_text_generation(payload) + task_params = await responses_request_to_text_generation(payload) resolved_model = await self._resolve_and_validate_text_model(task_params.model) task_params = task_params.model_copy(update={"model": resolved_model}) - command = TextGeneration(task_params=task_params) - await self._send(command) + command = await self._send_text_generation_with_images(task_params) if payload.stream: return StreamingResponse( @@ -1408,8 +1477,7 @@ class API: ) task_params = task_params.model_copy(update={"model": resolved_model}) - command = TextGeneration(task_params=task_params) - await self._send(command) + command = await self._send_text_generation_with_images(task_params) if payload.stream: return StreamingResponse( @@ -1445,8 +1513,7 @@ class API: ) task_params = task_params.model_copy(update={"model": resolved_model}) - command = TextGeneration(task_params=task_params) - await self._send(command) + command = await self._send_text_generation_with_images(task_params) if payload.stream: return StreamingResponse( diff --git a/src/exo/api/types/__init__.py b/src/exo/api/types/__init__.py index 4c6218d0..5821bdf4 100644 --- a/src/exo/api/types/__init__.py +++ b/src/exo/api/types/__init__.py @@ -6,7 +6,9 @@ from .api import BenchImageGenerationResponse as BenchImageGenerationResponse from .api import BenchImageGenerationTaskParams as BenchImageGenerationTaskParams from .api import CancelCommandResponse as CancelCommandResponse from .api import ChatCompletionChoice as ChatCompletionChoice +from .api import ChatCompletionContentPart as ChatCompletionContentPart from .api import ChatCompletionMessage as ChatCompletionMessage +from .api import ChatCompletionMessageImageUrl as ChatCompletionMessageImageUrl from .api import ChatCompletionMessageText as ChatCompletionMessageText from .api import ChatCompletionRequest as ChatCompletionRequest from .api import ChatCompletionResponse as ChatCompletionResponse diff --git a/src/exo/api/types/api.py b/src/exo/api/types/api.py index 828ee23b..b4fa3147 100644 --- a/src/exo/api/types/api.py +++ b/src/exo/api/types/api.py @@ -60,6 +60,14 @@ class ChatCompletionMessageText(BaseModel): text: str +class ChatCompletionMessageImageUrl(BaseModel): + type: Literal["image_url"] = "image_url" + image_url: dict[str, str] # {"url": "data:image/png;base64,..."} + + +ChatCompletionContentPart = ChatCompletionMessageText | ChatCompletionMessageImageUrl + + class ToolCallItem(BaseModel): id: str = Field(default_factory=lambda: str(uuid4())) name: str @@ -76,7 +84,7 @@ class ToolCall(BaseModel): class ChatCompletionMessage(BaseModel): role: Literal["system", "user", "assistant", "developer", "tool", "function"] content: ( - str | ChatCompletionMessageText | list[ChatCompletionMessageText] | None + str | ChatCompletionContentPart | list[ChatCompletionContentPart] | None ) = None reasoning_content: str | None = None name: str | None = None diff --git a/src/exo/api/types/claude_api.py b/src/exo/api/types/claude_api.py index 8eb18497..645da463 100644 --- a/src/exo/api/types/claude_api.py +++ b/src/exo/api/types/claude_api.py @@ -69,7 +69,7 @@ class ClaudeToolResultBlock(BaseModel, frozen=True): type: Literal["tool_result"] = "tool_result" tool_use_id: str - content: str | list[ClaudeTextBlock] | None = None + content: str | list[ClaudeTextBlock | ClaudeImageBlock] | None = None is_error: bool | None = None cache_control: dict[str, str] | None = None diff --git a/src/exo/api/types/ollama_api.py b/src/exo/api/types/ollama_api.py index 9f881fbe..58a54ac0 100644 --- a/src/exo/api/types/ollama_api.py +++ b/src/exo/api/types/ollama_api.py @@ -65,6 +65,7 @@ class OllamaGenerateRequest(BaseModel, frozen=True): keep_alive: str | int | None = None think: bool | None = None raw: bool = False + images: list[str] | None = None class OllamaGenerateResponse(BaseModel, frozen=True, strict=True): diff --git a/src/exo/api/types/openai_responses.py b/src/exo/api/types/openai_responses.py index eb21c4ae..b06593ac 100644 --- a/src/exo/api/types/openai_responses.py +++ b/src/exo/api/types/openai_responses.py @@ -27,6 +27,12 @@ class ResponseInputTextPart(BaseModel, frozen=True): text: str +class ResponseInputImagePart(BaseModel, frozen=True): + type: Literal["input_image"] = "input_image" + image_url: str | None = None + detail: str | None = None + + class ResponseOutputTextPart(BaseModel, frozen=True): """Output text content part (used when replaying assistant messages in input).""" @@ -34,7 +40,9 @@ class ResponseOutputTextPart(BaseModel, frozen=True): text: str -ResponseContentPart = ResponseInputTextPart | ResponseOutputTextPart +ResponseContentPart = ( + ResponseInputTextPart | ResponseInputImagePart | ResponseOutputTextPart +) # Request input item types diff --git a/src/exo/download/impl_shard_downloader.py b/src/exo/download/impl_shard_downloader.py index d87da8ee..0820380f 100644 --- a/src/exo/download/impl_shard_downloader.py +++ b/src/exo/download/impl_shard_downloader.py @@ -6,9 +6,18 @@ from typing import AsyncIterator, Callable from loguru import logger -from exo.download.download_utils import RepoDownloadProgress, download_shard +from exo.download.download_utils import ( + RepoDownloadProgress, + download_shard, +) from exo.download.shard_downloader import ShardDownloader -from exo.shared.models.model_cards import ModelCard, ModelId, get_model_cards +from exo.shared.models.model_cards import ( + ModelCard, + ModelId, + ModelTask, + get_model_cards, +) +from exo.shared.types.memory import Memory from exo.shared.types.worker.shards import ( PipelineShardMetadata, ShardMetadata, @@ -115,6 +124,38 @@ class ResumableShardDownloader(ShardDownloader): allow_patterns=allow_patterns, skip_internet=self.offline, ) + + if ( + not config_only + and not self.offline + and shard.model_card.vision + and shard.model_card.vision.weights_repo != str(shard.model_card.model_id) + ): + vision_repo = shard.model_card.vision.weights_repo + vision_card = ModelCard( + model_id=ModelId(vision_repo), + storage_size=Memory.from_bytes(0), + n_layers=1, + hidden_size=1, + supports_tensor=False, + tasks=[ModelTask.TextGeneration], + ) + vision_shard = PipelineShardMetadata( + model_card=vision_card, + device_rank=0, + world_size=1, + start_layer=0, + end_layer=1, + n_layers=1, + ) + await download_shard( + vision_shard, + self.on_progress_wrapper, + max_parallel_downloads=self.max_parallel_downloads, + allow_patterns=["*.safetensors", "config.json"], + skip_internet=self.offline, + ) + return target_dir async def get_shard_download_status( diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py index 7cc5a33c..0883d3e1 100644 --- a/src/exo/shared/models/model_cards.py +++ b/src/exo/shared/models/model_cards.py @@ -1,3 +1,4 @@ +import json from enum import Enum from typing import Annotated, Any @@ -13,6 +14,7 @@ from pydantic import ( Field, PositiveInt, ValidationError, + ValidationInfo, field_validator, model_validator, ) @@ -21,6 +23,7 @@ from tomlkit.exceptions import TOMLKitError from exo.shared.constants import ( EXO_CUSTOM_MODEL_CARDS_DIR, EXO_ENABLE_IMAGE_MODELS, + EXO_MODELS_DIRS, RESOURCES_DIR, ) from exo.shared.types.common import ModelId @@ -38,6 +41,23 @@ _BUILTIN_CARD_DIRS = [ _card_cache: dict[ModelId, "ModelCard"] = {} +def _detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None": + normalized = model_id.normalize() + for model_dir in [d / normalized for d in EXO_MODELS_DIRS]: + config_path = model_dir / "config.json" + if not config_path.exists(): + continue + try: + with open(config_path) as f: + raw = json.load(f) # type: ignore + return ConfigData.model_validate( + raw, context={"model_id": str(model_id)} + ).vision + except Exception: + continue + return None + + async def _load_cards_from_dir(directory: Path, *, is_custom: bool) -> None: """Load all TOML model cards from a directory into the cache.""" async for toml_file in directory.rglob("*.toml"): @@ -45,6 +65,10 @@ async def _load_cards_from_dir(directory: Path, *, is_custom: bool) -> None: card = await ModelCard.load_from_path(toml_file) if is_custom: card = card.model_copy(update={"is_custom": True}) + if card.vision is None: + vision = _detect_vision_from_config(card.model_id) + if vision is not None: + card = card.model_copy(update={"vision": vision}) if card.model_id not in _card_cache: _card_cache[card.model_id] = card except (ValidationError, TOMLKitError): @@ -89,6 +113,14 @@ class ComponentInfo(CamelCaseModel): safetensors_index_filename: str | None = None +class VisionCardConfig(CamelCaseModel): + image_token_id: int + model_type: str + weights_repo: str = "" + image_token: str | None = None + processor_repo: str | None = None + + class ModelCard(CamelCaseModel): model_id: ModelId storage_size: Memory @@ -105,6 +137,17 @@ class ModelCard(CamelCaseModel): uses_cfg: bool = False trust_remote_code: bool = True is_custom: bool = False + vision: VisionCardConfig | None = None + + @model_validator(mode="after") + def _fill_vision_weights_repo(self) -> "ModelCard": + if self.vision is not None and not self.vision.weights_repo: + object.__setattr__( + self, + "vision", + self.vision.model_copy(update={"weights_repo": str(self.model_id)}), + ) + return self @field_validator("tasks", mode="before") @classmethod @@ -162,6 +205,7 @@ class ModelCard(CamelCaseModel): tasks=[ModelTask.TextGeneration], trust_remote_code=False, is_custom=True, + vision=config_data.vision, ) @@ -196,6 +240,7 @@ class ConfigData(BaseModel): "decoder_layers", ) ) + vision: VisionCardConfig | None = None @property def supports_tensor(self) -> bool: @@ -217,24 +262,36 @@ class ConfigData(BaseModel): @model_validator(mode="before") @classmethod - def defer_to_text_config(cls, data: dict[str, Any]): + def defer_to_text_config(cls, data: dict[str, Any], info: ValidationInfo): text_config = data.get("text_config") - if text_config is None: - return data + if text_config is not None: + for field in [ + "architectures", + "hidden_size", + "num_key_value_heads", + "num_hidden_layers", + "num_layers", + "n_layer", + "n_layers", + "num_decoder_layers", + "decoder_layers", + ]: + if (val := text_config.get(field)) is not None: # pyright: ignore[reportAny] + data[field] = val - for field in [ - "architectures", - "hidden_size", - "num_key_value_heads", - "num_hidden_layers", - "num_layers", - "n_layer", - "n_layers", - "num_decoder_layers", - "decoder_layers", - ]: - if (val := text_config.get(field)) is not None: # pyright: ignore[reportAny] - data[field] = val + vision_config = data.get("vision_config") + image_token_id = data.get("image_token_id") + if vision_config is not None and image_token_id is not None: + model_type = str( + vision_config.get("model_type", data.get("model_type", "")) # pyright: ignore[reportAny] + ) + assert info.context is not None + + data["vision"] = VisionCardConfig( + image_token_id=int(image_token_id), # pyright: ignore[reportAny] + model_type=model_type, + weights_repo=info.context["model_id"], # type: ignore + ) return data @@ -257,7 +314,9 @@ async def fetch_config_data(model_id: ModelId) -> ConfigData: ), ) async with aiofiles.open(config_path, "r") as f: - return ConfigData.model_validate_json(await f.read()) + return ConfigData.model_validate_json( + await f.read(), context={"model_id": str(model_id)} + ) async def fetch_safetensors_size(model_id: ModelId) -> Memory: diff --git a/src/exo/shared/types/chunks.py b/src/exo/shared/types/chunks.py index 9faf0cd6..204556e8 100644 --- a/src/exo/shared/types/chunks.py +++ b/src/exo/shared/types/chunks.py @@ -68,6 +68,7 @@ class InputImageChunk(BaseChunk): data: str chunk_index: int total_chunks: int + image_index: int = 0 def __repr_args__(self) -> Generator[tuple[str, Any], None, None]: for name, value in super().__repr_args__(): # pyright: ignore[reportAny] diff --git a/src/exo/shared/types/text_generation.py b/src/exo/shared/types/text_generation.py index 3c2bcb48..29affa67 100644 --- a/src/exo/shared/types/text_generation.py +++ b/src/exo/shared/types/text_generation.py @@ -6,7 +6,7 @@ are converted to TextGenerationTaskParams at the API boundary via adapters. from typing import Any, Literal -from pydantic import BaseModel +from pydantic import BaseModel, Field from exo.shared.types.common import ModelId @@ -70,3 +70,7 @@ class TextGenerationTaskParams(BaseModel, frozen=True): min_p: float | None = None repetition_penalty: float | None = None repetition_context_size: int | None = None + images: list[str] = Field(default_factory=list) + image_hashes: dict[int, str] = Field(default_factory=dict) + total_input_chunks: int = 0 + image_count: int = 0 diff --git a/src/exo/worker/engines/mlx/cache.py b/src/exo/worker/engines/mlx/cache.py index 3877a61c..6376917d 100644 --- a/src/exo/worker/engines/mlx/cache.py +++ b/src/exo/worker/engines/mlx/cache.py @@ -1,5 +1,6 @@ import os from copy import deepcopy +from typing import TYPE_CHECKING import mlx.core as mx import psutil @@ -17,6 +18,9 @@ from exo.shared.types.mlx import KVCacheType, Model from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS from exo.worker.runner.bootstrap import logger +if TYPE_CHECKING: + from exo.worker.engines.mlx.vision import MediaRegion + # Fraction of device memory above which LRU eviction kicks in. # Smaller machines need more aggressive eviction. @@ -80,6 +84,7 @@ class KVPrefixCache: self.prompts: list[mx.array] = [] # mx array of tokens (ints) self.caches: list[KVCacheType] = [] self._snapshots: list[list[CacheSnapshot] | None] = [] + self._media_regions: list[list["MediaRegion"]] = [] self._last_used: list[int] = [] # monotonic counter of last access per entry self._access_counter: int = 0 self._group = group @@ -89,6 +94,7 @@ class KVPrefixCache: self.prompts.clear() self.caches.clear() self._snapshots.clear() + self._media_regions.clear() self._last_used.clear() def add_kv_cache( @@ -96,12 +102,14 @@ class KVPrefixCache: prompt_tokens: mx.array, cache: KVCacheType, ssm_snapshots: list[CacheSnapshot] | None = None, + media_regions: list["MediaRegion"] | None = None, ): """Add a new cache entry. Evicts LRU entries if memory is high.""" self._evict_if_needed() self.prompts.append(prompt_tokens) self.caches.append(deepcopy(cache)) self._snapshots.append(ssm_snapshots) + self._media_regions.append(media_regions or []) self._access_counter += 1 self._last_used.append(self._access_counter) logger.info(f"KV cache added: {len(prompt_tokens)} tokens") @@ -113,6 +121,7 @@ class KVPrefixCache: cache: KVCacheType, snapshots: list[CacheSnapshot] | None, restore_pos: int, + media_regions: list["MediaRegion"] | None = None, ): """Update an existing cache entry in-place.""" old_snapshots = self._snapshots[index] @@ -125,6 +134,7 @@ class KVPrefixCache: self.prompts[index] = prompt_tokens self.caches[index] = deepcopy(cache) self._snapshots[index] = merged or None + self._media_regions[index] = media_regions or [] self._access_counter += 1 self._last_used[index] = self._access_counter logger.info(f"KV cache updated (index {index}): {len(prompt_tokens)} tokens") @@ -149,6 +159,7 @@ class KVPrefixCache: self, model: Model, prompt_tokens: mx.array, + media_regions: list["MediaRegion"] | None = None, ) -> tuple[KVCacheType, mx.array, int | None]: """Get KV cache for prompt, returning remaining tokens to prefill. @@ -161,8 +172,13 @@ class KVPrefixCache: For models with SSM layers (which are ArraysCache in mlx), the cache is trimmed to the nearest SSM snapshot position at or before the match point for correctness. Same for rotating KV Cache. + + Media region validation: if the token-level prefix match extends into + a cached media region whose content_hash differs from the query's, the + match is truncated to the start of that region. """ max_length = len(prompt_tokens) + query_regions = media_regions or [] best_index: int | None = None best_length = 0 @@ -171,6 +187,12 @@ class KVPrefixCache: # Find best cache match for i, cached_prompt in enumerate(self.prompts): length = get_prefix_length(prompt_tokens, cached_prompt) + if length > 0: + length = self._validate_media_match( + length, + self._media_regions[i], + query_regions, + ) if length >= max_length - 1: best_index, best_length = i, length is_exact = True @@ -208,6 +230,37 @@ class KVPrefixCache: return prompt_cache, remaining, best_index + @staticmethod + def _validate_media_match( + match_length: int, + cached_regions: list["MediaRegion"], + query_regions: list["MediaRegion"], + ) -> int: + if not cached_regions: + return match_length + + query_by_start: dict[int, "MediaRegion"] = { + r.start_pos: r for r in query_regions + } + + for cached_r in cached_regions: + if cached_r.start_pos >= match_length: + break + query_r = query_by_start.get(cached_r.start_pos) + if query_r is None: + continue + if query_r.content_hash != cached_r.content_hash: + logger.info( + f"Media region mismatch at pos {cached_r.start_pos}: " + f"cached={cached_r.content_hash[:12]}... " + f"query={query_r.content_hash[:12]}... — " + f"truncating match from {match_length} to {cached_r.start_pos}" + ) + match_length = cached_r.start_pos + break + + return match_length + def _evict_if_needed(self): """Evict least recently used entries while memory usage is high.""" if len(self.caches) == 0: @@ -223,6 +276,7 @@ class KVPrefixCache: self.prompts.pop(lru_index) self.caches.pop(lru_index) self._snapshots.pop(lru_index) + self._media_regions.pop(lru_index) self._last_used.pop(lru_index) logger.info( f"KV cache evicted LRU entry ({evicted_tokens} tokens) due to memory usage" diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py index 3b026051..303609a4 100644 --- a/src/exo/worker/engines/mlx/generator/batch_generate.py +++ b/src/exo/worker/engines/mlx/generator/batch_generate.py @@ -1,3 +1,4 @@ +import contextlib import time from dataclasses import dataclass, field from typing import Callable, cast @@ -36,9 +37,16 @@ from exo.worker.engines.mlx.generator.generate import ( ban_token_ids, eos_ids_from_tokenizer, extract_top_logprobs, + patch_embed_tokens, prefill, ) from exo.worker.engines.mlx.utils_mlx import fix_unmatched_think_end_tokens +from exo.worker.engines.mlx.vision import ( + MediaRegion, + VisionProcessor, + VisionResult, + prepare_vision, +) from exo.worker.runner.bootstrap import logger _MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5 @@ -70,6 +78,7 @@ class _EngineTask: in_thinking: bool = False reasoning_tokens: int = 0 prefill_tps: float = 0.0 + media_regions: list[MediaRegion] = field(default_factory=list) @dataclass(eq=False) @@ -78,6 +87,7 @@ class ExoBatchGenerator: tokenizer: TokenizerWrapper group: mx.distributed.Group | None kv_prefix_cache: KVPrefixCache | None + vision_processor: VisionProcessor | None = None _mlx_gen: MlxBatchGenerator = field(init=False) _active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False) @@ -111,6 +121,27 @@ class ExoBatchGenerator: all_prompt_tokens, self.tokenizer ) + vision: VisionResult | None = None + media_regions: list[MediaRegion] = [] + + if self.vision_processor is not None: + try: + vision = prepare_vision( + images=task_params.images, + chat_template_messages=task_params.chat_template_messages, + vision_processor=self.vision_processor, + tokenizer=self.tokenizer, + model=self.model, + ) + except Exception: + logger.opt(exception=True).warning( + "Vision processing failed, falling back to text-only" + ) + + if vision is not None: + all_prompt_tokens = vision.prompt_tokens + media_regions = vision.media_regions + is_bench = task_params.bench prefix_hit_length = 0 @@ -119,7 +150,7 @@ class ExoBatchGenerator: if self.kv_prefix_cache is not None and not is_bench: cache, remaining_tokens, matched_index = self.kv_prefix_cache.get_kv_cache( - self.model, all_prompt_tokens + self.model, all_prompt_tokens, media_regions=media_regions ) prefix_hit_length = len(all_prompt_tokens) - len(remaining_tokens) if prefix_hit_length > 0: @@ -145,16 +176,24 @@ class ExoBatchGenerator: top_k=task_params.top_k if task_params.top_k is not None else 0, ) - _prefill_tps, _prefill_tokens, cache_snapshots = prefill( - self.model, - self.tokenizer, - sampler, - prompt_tokens[:-1], - cache, - self.group, - on_prefill_progress, - distributed_prompt_progress_callback, + vision_ctx = ( + patch_embed_tokens( + self.model, vision.embeddings, prefix_hit_length, len(prompt_tokens) - 1 + ) + if vision is not None + else contextlib.nullcontext() ) + with vision_ctx: + _prefill_tps, _prefill_tokens, cache_snapshots = prefill( + self.model, + self.tokenizer, + sampler, + prompt_tokens[:-1], + cache, + self.group, + on_prefill_progress, + distributed_prompt_progress_callback, + ) # We need to clamp rotating kv caches to max size so that mlx lm's _merge_caches behaves for c in cache: @@ -176,6 +215,7 @@ class ExoBatchGenerator: cache_snapshots, prefix_hit_length, matched_index, + media_regions, ) last_tokens = prompt_tokens[-2:] @@ -217,6 +257,7 @@ class ExoBatchGenerator: generation_start_time=time.perf_counter(), prefill_tps=_prefill_tps, generation_time_at_start=self._mlx_gen._stats.generation_time, + media_regions=media_regions, ) return uid @@ -383,6 +424,7 @@ class ExoBatchGenerator: cache_snapshots: list[CacheSnapshot] | None, prefix_hit_length: int, matched_index: int | None, + media_regions: list[MediaRegion] | None = None, ) -> None: if self.kv_prefix_cache is None: return @@ -402,10 +444,14 @@ class ExoBatchGenerator: cache, cache_snapshots, restore_pos=prefix_hit_length, + media_regions=media_regions, ) else: self.kv_prefix_cache.add_kv_cache( - all_prompt_tokens, cache, cache_snapshots + all_prompt_tokens, + cache, + cache_snapshots, + media_regions=media_regions, ) except Exception: logger.warning("Failed to save prefix cache", exc_info=True) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 19a9d455..d49ff58c 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -1,3 +1,4 @@ +import contextlib import functools import math import time @@ -55,6 +56,13 @@ from exo.worker.engines.mlx.utils_mlx import ( fix_unmatched_think_end_tokens, mx_barrier, ) +from exo.worker.engines.mlx.vision import ( + MediaRegion, + VisionProcessor, + VisionResult, + get_inner_model, + prepare_vision, +) from exo.worker.runner.bootstrap import logger generation_stream = mx.new_stream(mx.default_device()) @@ -62,6 +70,38 @@ generation_stream = mx.new_stream(mx.default_device()) _MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5 +@contextlib.contextmanager +def patch_embed_tokens( + model: Model, embeddings: mx.array, start_offset: int = 0, token_count: int = 0 +) -> Generator[None]: + inner = get_inner_model(model) # type: ignore + original_embed = inner.embed_tokens # type: ignore + end_offset = start_offset + token_count + offset = [start_offset] + + def _inject(input_ids: mx.array) -> mx.array: + start = offset[0] + if start >= end_offset: + return original_embed(input_ids) # type: ignore + chunk_len = input_ids.shape[-1] + end = min(start + chunk_len, end_offset) + offset[0] = end + if end - start < chunk_len: + return original_embed(input_ids) # type: ignore + return embeddings[:, start:end, :] + + for attr in dir(original_embed): # type: ignore + if not attr.startswith("_") and not hasattr(_inject, attr): + with contextlib.suppress(AttributeError, TypeError): + setattr(_inject, attr, getattr(original_embed, attr)) # type: ignore + + inner.embed_tokens = _inject + try: + yield + finally: + inner.embed_tokens = original_embed + + class PrefillCancelled(BaseException): """Raised when prefill is cancelled via the progress callback.""" @@ -447,6 +487,7 @@ def mlx_generate( on_prefill_progress: Callable[[int, int], None] | None = None, distributed_prompt_progress_callback: Callable[[], None] | None = None, on_generation_token: Callable[[], None] | None = None, + vision_processor: VisionProcessor | None = None, ) -> Generator[GenerationResponse]: # Ensure that generation stats only contains peak memory for this generation mx.reset_peak_memory() @@ -458,6 +499,24 @@ def mlx_generate( all_prompt_tokens = encode_prompt(tokenizer, prompt) all_prompt_tokens = fix_unmatched_think_end_tokens(all_prompt_tokens, tokenizer) + vision: VisionResult | None = None + if vision_processor is not None: + try: + vision = prepare_vision( + images=task.images, + chat_template_messages=task.chat_template_messages, + vision_processor=vision_processor, + tokenizer=tokenizer, + model=model, + ) + except Exception: + logger.opt(exception=True).warning( + "Vision processing failed, falling back to text-only" + ) + if vision is not None: + all_prompt_tokens = vision.prompt_tokens + media_regions: list[MediaRegion] = vision.media_regions if vision else [] + # Do not use the prefix cache if we are trying to do benchmarks. is_bench = task.bench if is_bench: @@ -471,7 +530,7 @@ def mlx_generate( prompt_tokens = all_prompt_tokens else: caches, prompt_tokens, matched_index = kv_prefix_cache.get_kv_cache( - model, all_prompt_tokens + model, all_prompt_tokens, media_regions=media_regions ) prefix_hit_length = len(all_prompt_tokens) - len(prompt_tokens) if prefix_hit_length > 0: @@ -505,17 +564,24 @@ def mlx_generate( ) max_stop_len = max((len(s) for s in stop_sequences), default=0) - # Prefill cache with all tokens except the last one - prefill_tps, prefill_tokens, ssm_snapshots_list = prefill( - model, - tokenizer, - sampler, - prompt_tokens[:-1], - caches, - group, - on_prefill_progress, - distributed_prompt_progress_callback, + maybe_vision_ctx = ( + patch_embed_tokens( + model, vision.embeddings, prefix_hit_length, len(prompt_tokens) - 1 + ) + if vision is not None + else contextlib.nullcontext() ) + with maybe_vision_ctx: + prefill_tps, prefill_tokens, ssm_snapshots_list = prefill( + model, + tokenizer, + sampler, + prompt_tokens[:-1], + caches, + group, + on_prefill_progress, + distributed_prompt_progress_callback, + ) cache_snapshots: list[CacheSnapshot] | None = ssm_snapshots_list or None # stream_generate starts from the last token @@ -655,10 +721,14 @@ def mlx_generate( caches, cache_snapshots, restore_pos=prefix_hit_length, + media_regions=media_regions, ) else: kv_prefix_cache.add_kv_cache( - full_prompt_tokens, caches, cache_snapshots + full_prompt_tokens, + caches, + cache_snapshots, + media_regions=media_regions, ) if on_generation_token is not None: diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index c8fc862c..4ee643b9 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -5,7 +5,10 @@ import sys import tempfile import time from pathlib import Path -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast + +if TYPE_CHECKING: + from exo.worker.engines.mlx.vision import VisionProcessor # Monkey-patch for transformers 5.x compatibility # Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location @@ -168,7 +171,7 @@ def load_mlx_items( group: Group | None, on_timeout: TimeoutCallback | None, on_layer_loaded: LayerLoadedCallback | None, -) -> tuple[Model, TokenizerWrapper]: +) -> "tuple[Model, TokenizerWrapper, VisionProcessor | None]": if group is None: logger.info(f"Single device used for {bound_instance.instance}") model_path = build_model_path(bound_instance.bound_shard.model_card.model_id) @@ -210,7 +213,18 @@ def load_mlx_items( mx.clear_cache() - return cast(Model, model), tokenizer + vision_config = bound_instance.bound_shard.model_card.vision + + if vision_config is not None: + from exo.worker.engines.mlx.vision import VisionProcessor + + vision_processor: VisionProcessor | None = VisionProcessor( + vision_config, bound_instance.bound_shard.model_card.model_id + ) + else: + vision_processor = None + + return cast(Model, model), tokenizer, vision_processor def shard_and_load( diff --git a/src/exo/worker/engines/mlx/vision.py b/src/exo/worker/engines/mlx/vision.py new file mode 100644 index 00000000..f0839b21 --- /dev/null +++ b/src/exo/worker/engines/mlx/vision.py @@ -0,0 +1,589 @@ +import base64 +import contextlib +import hashlib +import importlib +import inspect +import io +import json +import re +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from mlx_vlm.utils import ImageProcessor +import mlx.core as mx +import mlx.nn as nn +import numpy as np +from mlx_lm.tokenizer_utils import TokenizerWrapper +from mlx_vlm.prompt_utils import get_message_json +from mlx_vlm.utils import load_image_processor +from PIL import Image +from safetensors import safe_open +from transformers import AutoImageProcessor + +from exo.download.download_utils import build_model_path +from exo.shared.models.model_cards import VisionCardConfig +from exo.shared.types.common import ModelId +from exo.shared.types.mlx import Model +from exo.worker.engines.mlx.cache import encode_prompt +from exo.worker.engines.mlx.utils_mlx import fix_unmatched_think_end_tokens +from exo.worker.runner.bootstrap import logger + + +def _filter_config(cls: type, d: dict[str, Any]) -> dict[str, Any]: + valid = set(inspect.signature(cls.__init__).parameters.keys()) - {"self"} + return {k: v for k, v in d.items() if k in valid} # type: ignore + + +_video_processor_patched = False + + +def _patch_video_processor() -> None: + """Patch so we don't crash horribly when torch vision isn't installed""" + # TODO: Update if we add torch vision. + global _video_processor_patched + if _video_processor_patched: + return + try: + from transformers.processing_utils import MODALITY_TO_AUTOPROCESSOR_MAPPING + + mapping = MODALITY_TO_AUTOPROCESSOR_MAPPING._MAPPING_NAMES # type: ignore + mapping.pop("video_processor", None) + except (ImportError, AttributeError): + pass + _video_processor_patched = True + + +def decode_base64_image(b64_data: str) -> Image.Image: + raw = base64.b64decode(b64_data) + img = Image.open(io.BytesIO(raw)) + return img.convert("RGB") + + +def _format_vlm_messages( + messages: list[dict[str, Any]], + model_type: str, +) -> list[dict[str, Any]]: + formatted: list[dict[str, Any]] = [] + for msg in messages: + role: str = str(msg.get("role", "user")) # type: ignore + content: Any = msg.get("content") + if not isinstance(content, list): + formatted.append(msg) + continue + parts: list[dict[str, Any]] = content # type: ignore + text_parts = [str(p["text"]) for p in parts if p.get("type") == "text"] # type: ignore + n_images = sum(1 for p in parts if p.get("type") in ("image", "image_url")) + result: dict[str, Any] = get_message_json( + model_type, " ".join(text_parts), role, num_images=n_images + ) + formatted.append(result) + return formatted + + +def build_vision_prompt( + tokenizer: TokenizerWrapper, + chat_template_messages: list[dict[str, Any]], + n_tokens_per_image: list[int], + image_token: str, +) -> str: + logger.info( + f"Vision prompt messages: {[{k: (v[:50] if isinstance(v, str) else v) for k, v in m.items()} for m in chat_template_messages]}" # type: ignore + ) + prompt: str = tokenizer.apply_chat_template( + chat_template_messages, + tokenize=False, + add_generation_prompt=True, + ) + + image_idx = 0 + result: list[str] = [] + i = 0 + pad_len = len(image_token) + while i < len(prompt): + if prompt[i : i + pad_len] == image_token: + n = ( + n_tokens_per_image[image_idx] + if image_idx < len(n_tokens_per_image) + else 1 + ) + result.append(image_token * n) + image_idx += 1 + i += pad_len + else: + result.append(prompt[i]) + i += 1 + + return "".join(result) + + +@dataclass +class MediaRegion: + content_hash: str + start_pos: int + end_pos: int + + +@dataclass +class VisionResult: + prompt: str + prompt_tokens: mx.array + embeddings: mx.array + media_regions: list[MediaRegion] + + +class VisionEncoder: + def __init__(self, config: VisionCardConfig, model_id: ModelId): + self._config = config + self._main_model_path = build_model_path(model_id) + self._model_path = build_model_path(ModelId(config.weights_repo)) + self._vision_tower: nn.Module | None = None + self._projector: nn.Module | None = None + self._processor: "ImageProcessor | None" = None + self._spatial_merge_size: int = 2 + self._merge_kernel_size: list[int] | None = None + self._needs_nhwc: bool = False + self._loaded = False + + def _load_config_json(self) -> dict[str, Any]: + for candidate in (self._main_model_path, self._model_path): + path = candidate / "config.json" + if path.exists(): + with open(path) as f: + return json.load(f) # type: ignore + return {} + + def _import_mlx_vlm(self, *submodules: str) -> Any: # type: ignore + mt = self._config.model_type + results: list[Any] = [] + for sub in submodules: + name = f"mlx_vlm.models.{mt}.{sub}" + results.append(importlib.import_module(name)) + return results[0] if len(results) == 1 else tuple(results) + + def ensure_loaded(self) -> None: + if self._loaded: + return + self._load_weights() + self._loaded = True + + def _load_weights(self) -> None: + _patch_video_processor() + logger.info(f"Loading vision weights from {self._model_path}") + config = self._load_config_json() + if not config: + raise FileNotFoundError(f"config.json not found in {self._model_path}") + + vision_cfg = config.get("vision_config", {}) # type: ignore + + config_mod, vision_mod = self._import_mlx_vlm("config", "vision") # type: ignore + vision_config_cls = config_mod.VisionConfig # type: ignore + vision_model_cls = vision_mod.VisionModel # type: ignore + + vision_config = vision_config_cls( # type: ignore + **_filter_config(vision_config_cls, vision_cfg) # type: ignore + ) + self._spatial_merge_size = getattr(vision_config, "spatial_merge_size", 2) # type: ignore + self._vision_tower = vision_model_cls(vision_config) + model_mod: Any = None + with contextlib.suppress(ImportError): + model_mod = self._import_mlx_vlm(self._config.model_type) # type: ignore + + projector_cls = None + if model_mod is not None: + for attr_name in dir(model_mod): # type: ignore + obj = getattr(model_mod, attr_name) # type: ignore + if ( + isinstance(obj, type) + and issubclass(obj, nn.Module) + and "Projector" in attr_name + ): + projector_cls = obj + break + + if projector_cls is not None: + text_config = config_mod.TextConfig( # type: ignore + **_filter_config(config_mod.TextConfig, config.get("text_config", {})) # type: ignore + ) + extra = { + k: v + for k, v in config.items() # type: ignore + if k not in ("text_config", "vision_config") + } + extra.setdefault("model_type", self._config.model_type) + model_config = config_mod.ModelConfig( # type: ignore + text_config=text_config, + vision_config=vision_config, + **_filter_config(config_mod.ModelConfig, extra), # type: ignore + ) + self._projector = projector_cls(model_config) # type: ignore + + processor_repo = self._config.processor_repo + if processor_repo: + self._load_weights_from_separate_repo() + else: + self._load_weights_from_model_repo() + + repo = processor_repo or str(self._model_path) + image_proc = load_image_processor(repo) + if image_proc is not None: + self._processor = image_proc + else: + self._processor = AutoImageProcessor.from_pretrained( # type: ignore + repo, trust_remote_code=True + ) + if processor_repo: + self._merge_kernel_size = vision_cfg.get("merge_kernel_size", [2, 2]) # type: ignore + self._needs_nhwc = True + logger.info(f"HF image processor loaded from {repo}") + + def _load_weights_from_separate_repo(self) -> None: + safetensors_files = list(self._model_path.glob("*.safetensors")) + if not safetensors_files: + raise FileNotFoundError(f"No safetensors files found in {self._model_path}") + + weights: dict[str, mx.array] = {} + for sf_path in safetensors_files: + with safe_open(str(sf_path), framework="pt") as f: + keys = f.keys() + for key in keys: + tensor = f.get_tensor(key) # type: ignore + np_tensor = tensor.float().numpy() # type: ignore + weights[key] = mx.array(np_tensor, dtype=mx.bfloat16) # type: ignore + + vision_weights: dict[str, mx.array] = {} + projector_weights: dict[str, mx.array] = {} + for key, val in weights.items(): + if key.startswith("vision_tower."): + short_key = key[len("vision_tower.") :] + if short_key.startswith("encoder."): + short_key = short_key[len("encoder.") :] + m = re.match(r"^(blocks\.\d+)\.(wqkv|wo)\.(weight|bias)$", short_key) + if m: + short_key = f"{m.group(1)}.attn.{m.group(2)}.{m.group(3)}" + if short_key == "patch_embed.proj.weight" and val.ndim == 4: + val = val.transpose(0, 2, 3, 1) + vision_weights[short_key] = val + elif key.startswith(("mm_projector.", "multi_modal_projector.")): + if key.startswith("multi_modal_projector."): + short_key = key[len("multi_modal_projector.") :] + if short_key.startswith("mm_projector."): + short_key = short_key[len("mm_projector.") :] + else: + short_key = key[len("mm_projector.") :] + short_key = short_key.replace("proj.0.", "linear_1.").replace( + "proj.2.", "linear_2." + ) + projector_weights[short_key] = val + + assert self._vision_tower is not None + self._vision_tower.load_weights(list(vision_weights.items())) + mx.eval(self._vision_tower.parameters()) + + if self._projector is not None and projector_weights: + self._projector.load_weights(list(projector_weights.items())) + mx.eval(self._projector.parameters()) + + n_vision = sum(v.size for _, v in vision_weights.items()) + n_proj = sum(v.size for _, v in projector_weights.items()) + logger.info( + f"Vision encoder loaded: {n_vision / 1e6:.1f}M params" + + (f", projector: {n_proj / 1e6:.1f}M params" if n_proj else "") + ) + + def _load_weights_from_model_repo(self) -> None: + safetensors_files = sorted(self._model_path.glob("*.safetensors")) + if not safetensors_files: + raise FileNotFoundError(f"No safetensors files found in {self._model_path}") + + vision_prefixes = ["vision_tower.", "model.visual."] + vision_weights: dict[str, mx.array] = {} + found_raw_prefix = False + for sf_path in safetensors_files: + file_weights: dict[str, mx.array] = mx.load(str(sf_path)) # type: ignore + for key, val in file_weights.items(): + for prefix in vision_prefixes: + if key.startswith(prefix): + short_key = key[len(prefix) :] + vision_weights[short_key] = val + if prefix == "model.visual.": + found_raw_prefix = True + break + + if not vision_weights: + raise ValueError( + f"No vision weights found with prefixes {vision_prefixes} in {self._model_path}. " + "Ensure the model repo contains bundled vision weights." + ) + + assert self._vision_tower is not None + if found_raw_prefix and hasattr(self._vision_tower, "sanitize"): + vision_weights = self._vision_tower.sanitize(vision_weights) # type: ignore + + self._vision_tower.load_weights(list(vision_weights.items())) # type: ignore + mx.eval(self._vision_tower.parameters()) + + n_vision = sum(v.size for _, v in vision_weights.items()) # type: ignore + logger.info(f"Vision encoder loaded: {n_vision / 1e6:.1f}M params") + + def encode_images(self, images: list[str]) -> tuple[mx.array, list[int]]: + self.ensure_loaded() + assert self._vision_tower is not None + assert self._processor is not None + + pil_images = [decode_base64_image(b64) for b64 in images] + for idx, img in enumerate(pil_images): + logger.info(f"Image {idx}: {img.width}x{img.height} mode={img.mode}") + + if self._config.processor_repo: + processed = self._processor.preprocess( + [{"type": "image", "image": img} for img in pil_images], + return_tensors="np", + ) + pixel_values = mx.array(processed["pixel_values"]) # type: ignore + grid_thw = mx.array(processed["grid_thws"]) # type: ignore + assert self._merge_kernel_size is not None + merge_length = int(np.prod(self._merge_kernel_size)) + n_tokens_per_image = [ + int(mx.prod(grid_thw[i]).item()) // merge_length + for i in range(grid_thw.shape[0]) + ] + else: + processed = self._processor( + images=pil_images, + return_tensors="np", + ) + pixel_values = mx.array(processed["pixel_values"]) # type: ignore + grid_thw = mx.array(processed["image_grid_thw"]) # type: ignore + merge_unit = self._spatial_merge_size**2 + n_tokens_per_image = [ + int( + grid_thw[i, 0].item() + * grid_thw[i, 1].item() + * grid_thw[i, 2].item() + ) + // merge_unit + for i in range(grid_thw.shape[0]) + ] + + if self._needs_nhwc: + grid_hw = grid_thw[:, 1:] if grid_thw.shape[-1] == 3 else grid_thw + hidden_states = self._vision_tower( + pixel_values.transpose(0, 2, 3, 1), + output_hidden_states=True, + grid_thw=grid_hw, + ) + else: + result = self._vision_tower(pixel_values, grid_thw) + hidden_states = result[0] if isinstance(result, tuple) else result + + if self._projector is not None: + image_features: mx.array = self._projector(hidden_states) + else: + image_features = hidden_states + + return image_features, n_tokens_per_image + + +def get_inner_model(model: nn.Module) -> Any: # type: ignore + for candidate in ( + getattr(model, "model", None), + getattr(getattr(model, "language_model", None), "model", None), + ): + if candidate is not None and hasattr(candidate, "embed_tokens"): # type: ignore + return candidate # type: ignore + + raise ValueError( + f"Could not find inner transformer (embed_tokens) in {type(model).__name__}. " + "Add a new pattern to _get_inner_model() for this architecture." + ) + + +def create_vision_embeddings( + model: Model, + prompt_tokens: mx.array, + image_features: mx.array, + image_token_id: int, +) -> mx.array: + inner = get_inner_model(model) # type: ignore + embed_tokens = inner.embed_tokens # type: ignore + + input_embeddings: mx.array = embed_tokens(prompt_tokens[None]) # type: ignore + + is_image: mx.array = mx.equal(prompt_tokens, image_token_id) + n_placeholders = int(mx.sum(is_image).item()) + + if n_placeholders > 0: + if n_placeholders != image_features.shape[0]: + logger.warning( + f"Placeholder count ({n_placeholders}) != image features " + f"({image_features.shape[0]}). Using min of both." + ) + n = min(n_placeholders, image_features.shape[0]) + image_features = image_features[:n] + + image_indices = mx.cumsum(is_image.astype(mx.int32)) - 1 + image_indices = mx.clip(image_indices, 0, image_features.shape[0] - 1) + + gathered = image_features[image_indices].astype(input_embeddings.dtype) + result = mx.where(is_image[:, None], gathered, input_embeddings[0]) + input_embeddings = result[None] + + return input_embeddings + + +def _find_media_regions( + prompt_tokens: mx.array, + images: list[str], + image_token_id: int, +) -> list[MediaRegion]: + tokens_np = np.array(prompt_tokens) + is_pad = tokens_np == image_token_id # type: ignore + + regions: list[MediaRegion] = [] + in_run = False + run_start = 0 + for pos, pad in enumerate(is_pad): # type: ignore + if pad and not in_run: + run_start = pos + in_run = True + elif not pad and in_run: + regions.append( + MediaRegion(content_hash="", start_pos=run_start, end_pos=pos) + ) + in_run = False + if in_run: + regions.append( + MediaRegion(content_hash="", start_pos=run_start, end_pos=len(tokens_np)) + ) + + for i, region in enumerate(regions): + if i < len(images): + img = decode_base64_image(images[i]) + region.content_hash = hashlib.sha256(img.tobytes()).hexdigest() + else: + logger.warning(f"Media region {i} has no corresponding image") + + return regions + + +class VisionProcessor: + """ + Pipeline for vision models: + 1. Encode images into features (or grab from cache) + 2. Replace image placeholders with the features + 3. Build vision prompt + 4. Provide media regions for prefix caching + """ + + def __init__(self, config: VisionCardConfig, model_id: ModelId): + self.vision_config = config + self._encoder = VisionEncoder(config, model_id) + self._feature_cache: dict[str, tuple[mx.array, list[int]]] = {} + self._feature_cache_max = 32 + + def load(self) -> None: + self._encoder.ensure_loaded() + + def _image_cache_key(self, images: list[str]) -> str: + h = hashlib.sha256() + for img in images: + pil = decode_base64_image(img) + h.update(pil.tobytes()) + return h.hexdigest() + + def process( + self, + images: list[str], + chat_template_messages: list[dict[str, Any]], + tokenizer: TokenizerWrapper, + model: Model, + ) -> VisionResult: + logger.info(f"Vision pipeline: {len(images)} image(s)") + + cache_key = self._image_cache_key(images) + cached = self._feature_cache.pop(cache_key, None) + if cached is not None: + self._feature_cache[cache_key] = cached + image_features, n_tokens_per_image = cached + else: + image_features, n_tokens_per_image = self._encoder.encode_images(images) + self._feature_cache[cache_key] = (image_features, n_tokens_per_image) + while len(self._feature_cache) > self._feature_cache_max: + del self._feature_cache[next(iter(self._feature_cache))] + logger.info( + f"Vision features: {image_features.shape} " + f"({image_features.shape[0]} tokens, per-image: {n_tokens_per_image})" + ) + + image_token = self.vision_config.image_token + if image_token is None: + image_token = tokenizer.decode([self.vision_config.image_token_id]) + + formatted_messages = _format_vlm_messages( + chat_template_messages, self.vision_config.model_type + ) + + prompt = build_vision_prompt( + tokenizer, + formatted_messages, + n_tokens_per_image, + image_token, + ) + + logger.info( + f"Expanded prompt has {prompt.count(image_token)} image_token occurrences, total len={len(prompt)}" + ) + + prompt_tokens: mx.array = encode_prompt(tokenizer, prompt) + prompt_tokens = fix_unmatched_think_end_tokens(prompt_tokens, tokenizer) + n_image_tokens = int( + mx.sum(mx.equal(prompt_tokens, self.vision_config.image_token_id)).item() + ) + logger.info( + f"Encoded prompt: {len(prompt_tokens)} tokens, {n_image_tokens} image pad tokens" + ) + + embeddings = create_vision_embeddings( + model, + prompt_tokens, + image_features, + self.vision_config.image_token_id, + ) + mx.eval(embeddings) + + media_regions = _find_media_regions( + prompt_tokens, + images, + self.vision_config.image_token_id, + ) + + return VisionResult( + prompt=prompt, + prompt_tokens=prompt_tokens, + embeddings=embeddings, + media_regions=media_regions, + ) + + +def prepare_vision( + images: list[str] | None, + chat_template_messages: list[dict[str, Any]] | None, + vision_processor: VisionProcessor, + tokenizer: TokenizerWrapper, + model: Model, +) -> VisionResult | None: + if not images: + return None + if chat_template_messages is None: + logger.warning( + "Vision request missing chat_template_messages — ignoring images" + ) + return None + + return vision_processor.process( + images=images, + chat_template_messages=chat_template_messages, + tokenizer=tokenizer, + model=model, + ) diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index 30128594..ded0481b 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -1,3 +1,4 @@ +import hashlib from collections import defaultdict from datetime import datetime, timezone @@ -9,6 +10,7 @@ from exo.api.types import ImageEditsTaskParams from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model from exo.shared.apply import apply from exo.shared.models.model_cards import ModelId, add_to_card_cache, delete_custom_card +from exo.shared.types.chunks import InputImageChunk from exo.shared.types.commands import ( ForwarderCommand, ForwarderDownloadCommand, @@ -38,6 +40,7 @@ from exo.shared.types.tasks import ( Shutdown, Task, TaskStatus, + TextGeneration, ) from exo.shared.types.topology import Connection, SocketConnection from exo.shared.types.worker.downloads import DownloadCompleted @@ -76,8 +79,9 @@ class Worker: self._system_id = SystemId() # Buffer for input image chunks (for image editing) - self.input_chunk_buffer: dict[CommandId, dict[int, str]] = {} + self.input_chunk_buffer: dict[CommandId, dict[int, InputImageChunk]] = {} self.input_chunk_counts: dict[CommandId, int] = {} + self.image_cache: dict[str, str] = {} self._download_backoff: KeyedBackoff[ModelId] = KeyedBackoff(base=0.5, cap=10.0) self._stopped: anyio.Event = anyio.Event() @@ -131,7 +135,7 @@ class Worker: self.input_chunk_counts[cmd_id] = event.chunk.total_chunks self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = ( - event.chunk.data + event.chunk ) if isinstance(event, CustomModelCardAdded): @@ -152,7 +156,6 @@ class Worker: self.state.runners, self.state.tasks, self.input_chunk_buffer, - self.input_chunk_counts, ) if task is None: continue @@ -245,7 +248,7 @@ class Worker: # Assemble image from chunks and inject into task cmd_id = task.command_id chunks = self.input_chunk_buffer.get(cmd_id, {}) - assembled = "".join(chunks[i] for i in range(len(chunks))) + assembled = "".join(chunks[i].data for i in range(len(chunks))) logger.info( f"Assembled input image from {len(chunks)} chunks, " f"total size: {len(assembled)} bytes" @@ -279,6 +282,52 @@ class Worker: if cmd_id in self.input_chunk_counts: del self.input_chunk_counts[cmd_id] await self._start_runner_task(modified_task) + + case TextGeneration() if ( + task.task_params.image_hashes + or task.task_params.total_input_chunks > 0 + ): + cmd_id = task.command_id + by_index: dict[int, str] = {} + + for idx, h in task.task_params.image_hashes.items(): + assert h in self.image_cache + by_index[idx] = self.image_cache[h] + + if task.task_params.total_input_chunks > 0: + chunk_buffer = self.input_chunk_buffer.get(cmd_id, {}) + per_image: defaultdict[int, list[InputImageChunk]] = ( + defaultdict(list) + ) + for chunk in chunk_buffer.values(): + per_image[chunk.image_index].append(chunk) + for img_idx in sorted(per_image): + sorted_chunks = sorted( + per_image[img_idx], key=lambda c: c.chunk_index + ) + img = "".join(c.data for c in sorted_chunks) + self.image_cache[ + hashlib.sha256(img.encode("ascii")).hexdigest() + ] = img + by_index[img_idx] = img + logger.info( + f"Assembled {len(per_image)} VLM image(s) " + f"from {len(chunk_buffer)} chunks" + ) + + resolved_images = [by_index[i] for i in sorted(by_index)] + modified_task = task.model_copy( + update={ + "task_params": task.task_params.model_copy( + update={"images": resolved_images} + ) + } + ) + if cmd_id in self.input_chunk_buffer: + del self.input_chunk_buffer[cmd_id] + if cmd_id in self.input_chunk_counts: + del self.input_chunk_counts[cmd_id] + await self._start_runner_task(modified_task) case task: await self._start_runner_task(task) diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py index 9b268ae6..0df10ca4 100644 --- a/src/exo/worker/plan.py +++ b/src/exo/worker/plan.py @@ -2,6 +2,7 @@ from collections.abc import Mapping, Sequence +from exo.shared.types.chunks import InputImageChunk from exo.shared.types.common import CommandId, NodeId from exo.shared.types.tasks import ( CancelTask, @@ -49,8 +50,7 @@ def plan( instances: Mapping[InstanceId, Instance], all_runners: Mapping[RunnerId, RunnerStatus], # all global tasks: Mapping[TaskId, Task], - input_chunk_buffer: Mapping[CommandId, dict[int, str]] | None = None, - input_chunk_counts: Mapping[CommandId, int] | None = None, + input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]] | None = None, ) -> Task | None: # Python short circuiting OR logic should evaluate these sequentially. return ( @@ -272,7 +272,7 @@ def _pending_tasks( runners: Mapping[RunnerId, RunnerSupervisor], tasks: Mapping[TaskId, Task], all_runners: Mapping[RunnerId, RunnerStatus], - input_chunk_buffer: Mapping[CommandId, dict[int, str]], + input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]] | None, ) -> Task | None: for task in tasks.values(): # for now, just forward chat completions @@ -282,12 +282,15 @@ def _pending_tasks( if task.task_status not in (TaskStatus.Pending, TaskStatus.Running): continue - # For ImageEdits tasks, verify all input chunks have been received - if isinstance(task, ImageEdits) and task.task_params.total_input_chunks > 0: + # For tasks with images, verify all input chunks have been received + expected_image_chunks = 0 + if isinstance(task, (ImageEdits, TextGeneration)): + expected_image_chunks = task.task_params.total_input_chunks + if expected_image_chunks > 0: + assert input_chunk_buffer is not None cmd_id = task.command_id - expected = task.task_params.total_input_chunks received = len(input_chunk_buffer.get(cmd_id, {})) - if received < expected: + if received < expected_image_chunks: continue # Wait for all chunks to arrive for runner in runners.values(): diff --git a/src/exo/worker/runner/llm_inference/batch_generator.py b/src/exo/worker/runner/llm_inference/batch_generator.py index 5998ef08..5284d53a 100644 --- a/src/exo/worker/runner/llm_inference/batch_generator.py +++ b/src/exo/worker/runner/llm_inference/batch_generator.py @@ -29,6 +29,7 @@ from exo.worker.engines.mlx.utils_mlx import ( mx_all_gather_tasks, mx_any, ) +from exo.worker.engines.mlx.vision import VisionProcessor from exo.worker.runner.bootstrap import logger from .model_output_parsers import apply_all_parsers @@ -120,6 +121,7 @@ class SequentialGenerator(InferenceGenerator): device_rank: int cancel_receiver: MpReceiver[TaskId] event_sender: MpSender[Event] + vision_processor: VisionProcessor | None = None check_for_cancel_every: int = 50 _cancelled_tasks: set[TaskId] = field(default_factory=set, init=False) @@ -304,6 +306,7 @@ class SequentialGenerator(InferenceGenerator): distributed_prompt_progress_callback=distributed_prompt_progress_callback, on_generation_token=on_generation_token, group=self.group, + vision_processor=self.vision_processor, ) def close(self) -> None: @@ -322,6 +325,7 @@ class BatchGenerator(InferenceGenerator): cancel_receiver: MpReceiver[TaskId] event_sender: MpSender[Event] check_for_cancel_every: int = 50 + vision_processor: VisionProcessor | None = None _cancelled_tasks: set[TaskId] = field(default_factory=set, init=False) _maybe_queue: list[TextGeneration] = field(default_factory=list, init=False) @@ -344,6 +348,7 @@ class BatchGenerator(InferenceGenerator): tokenizer=self.tokenizer, group=self.group, kv_prefix_cache=self.kv_prefix_cache, + vision_processor=self.vision_processor, ) def warmup(self): diff --git a/src/exo/worker/runner/llm_inference/runner.py b/src/exo/worker/runner/llm_inference/runner.py index 4f2f3d30..5c19ec73 100644 --- a/src/exo/worker/runner/llm_inference/runner.py +++ b/src/exo/worker/runner/llm_inference/runner.py @@ -57,6 +57,7 @@ from exo.worker.engines.mlx.utils_mlx import ( initialize_mlx, load_mlx_items, ) +from exo.worker.engines.mlx.vision import VisionProcessor from exo.worker.runner.bootstrap import logger from exo.worker.runner.llm_inference.batch_generator import ( BatchGenerator, @@ -103,7 +104,9 @@ class Runner: self.setup_start_time = time.time() self.generator: Builder | InferenceGenerator = Builder( - self.model_id, self.event_sender, self.cancel_receiver + self.model_id, + self.event_sender, + self.cancel_receiver, ) self.seen: set[TaskId] = set() @@ -195,13 +198,15 @@ class Runner: assert ( ModelTask.TextGeneration in self.shard_metadata.model_card.tasks ), f"Incorrect model task(s): {self.shard_metadata.model_card.tasks}" - self.generator.inference_model, self.generator.tokenizer = ( - load_mlx_items( - self.bound_instance, - self.generator.group, - on_timeout=on_model_load_timeout, - on_layer_loaded=on_layer_loaded, - ) + ( + self.generator.inference_model, + self.generator.tokenizer, + self.generator.vision_processor, + ) = load_mlx_items( + self.bound_instance, + self.generator.group, + on_timeout=on_model_load_timeout, + on_layer_loaded=on_layer_loaded, ) self.generator = self.generator.build() @@ -381,6 +386,7 @@ class Builder: inference_model: Model | None = None tokenizer: TokenizerWrapper | None = None group: mx.distributed.Group | None = None + vision_processor: VisionProcessor | None = None def build( self, @@ -389,6 +395,8 @@ class Builder: assert self.inference_model assert self.tokenizer + vision_processor = self.vision_processor + tool_parser = None logger.info( f"model has_tool_calling={self.tokenizer.has_tool_calling} using tokens {self.tokenizer.tool_call_start}, {self.tokenizer.tool_call_end}" @@ -419,6 +427,7 @@ class Builder: device_rank=device_rank, cancel_receiver=self.cancel_receiver, event_sender=self.event_sender, + vision_processor=vision_processor, ) logger.info("using BatchGenerator") return BatchGenerator( @@ -431,4 +440,5 @@ class Builder: device_rank=device_rank, cancel_receiver=self.cancel_receiver, event_sender=self.event_sender, + vision_processor=vision_processor, ) diff --git a/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py b/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py index d6e3774b..2f4ca7e6 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py @@ -137,7 +137,9 @@ async def test_tokenizer_encode_decode(model_card: ModelCard) -> None: # Test decoding decoded = tokenizer.decode(encoded) assert isinstance(decoded, str), f"decode() should return a string for {model_id}" - assert test_text in decoded or decoded.strip() == test_text.strip(), ( + normalized_decoded = decoded.replace(" ", "").lower() + normalized_expected = test_text.replace(" ", "").lower() + assert normalized_expected in normalized_decoded, ( f"decode(encode(x)) should preserve text for {model_id}: got {decoded!r}" ) diff --git a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py index 8c651cde..cd4c18d7 100644 --- a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py +++ b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py @@ -115,7 +115,9 @@ def assert_events_equal(test_events: Iterable[Event], true_events: Iterable[Even def patch_out_mlx(monkeypatch: pytest.MonkeyPatch): # initialize_mlx returns a mock group monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MockGroup())) - monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((1, MockTokenizer))) + monkeypatch.setattr( + mlx_runner, "load_mlx_items", make_nothin((1, MockTokenizer, None)) + ) monkeypatch.setattr(mlx_batch_generator, "warmup_inference", make_nothin(1)) monkeypatch.setattr(mlx_batch_generator, "_check_for_debug_prompts", nothin) monkeypatch.setattr(mlx_batch_generator, "mx_any", make_nothin(False)) diff --git a/tests/test_vision_cache.py b/tests/test_vision_cache.py new file mode 100644 index 00000000..46b47bd8 --- /dev/null +++ b/tests/test_vision_cache.py @@ -0,0 +1,63 @@ +from exo.worker.engines.mlx.cache import KVPrefixCache +from exo.worker.engines.mlx.vision import MediaRegion + +validate = KVPrefixCache._validate_media_match + + +class TestValidateMediaMatch: + def test_text_only_no_truncation(self): + assert validate(8000, [], []) == 8000 + + def test_text_prefix_before_image(self): + cached = [MediaRegion("hashA", 5000, 8600)] + assert validate(5000, cached, []) == 5000 + + def test_same_image_same_position(self): + cached = [MediaRegion("hashA", 5000, 8600)] + query = [MediaRegion("hashA", 5000, 8600)] + assert validate(9000, cached, query) == 9000 + + def test_different_image_truncates(self): + cached = [MediaRegion("hashA", 5000, 8600)] + query = [MediaRegion("hashB", 5000, 8600)] + assert validate(9000, cached, query) == 5000 + + def test_match_below_region_start(self): + cached = [MediaRegion("hashA", 5000, 8600)] + query = [MediaRegion("hashB", 5000, 8600)] + assert validate(4000, cached, query) == 4000 + + def test_text_followup_no_images_in_query(self): + cached = [MediaRegion("hashA", 5000, 8600)] + assert validate(9000, cached, []) == 9000 + + def test_multiple_images_first_mismatch_truncates(self): + cached = [ + MediaRegion("hashA", 2000, 4000), + MediaRegion("hashB", 6000, 8000), + ] + query = [ + MediaRegion("hashA", 2000, 4000), + MediaRegion("hashC", 6000, 8000), + ] + assert validate(9000, cached, query) == 6000 + + def test_multiple_images_all_match(self): + cached = [ + MediaRegion("hashA", 2000, 4000), + MediaRegion("hashB", 6000, 8000), + ] + query = [ + MediaRegion("hashA", 2000, 4000), + MediaRegion("hashB", 6000, 8000), + ] + assert validate(9000, cached, query) == 9000 + + def test_no_cached_regions(self): + query = [MediaRegion("hashA", 100, 200)] + assert validate(500, [], query) == 500 + + def test_cached_region_beyond_match(self): + cached = [MediaRegion("hashA", 10000, 12000)] + query = [MediaRegion("hashB", 10000, 12000)] + assert validate(5000, cached, query) == 5000 diff --git a/uv.lock b/uv.lock index 892a2503..7ddd1486 100644 --- a/uv.lock +++ b/uv.lock @@ -207,20 +207,32 @@ name = "cffi" version = "2.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pycparser", marker = "implementation_name != 'PyPy' and sys_platform == 'linux'" }, + { name = "pycparser", marker = "(implementation_name != 'PyPy' and sys_platform == 'darwin') or (implementation_name != 'PyPy' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8453301356628e8147c79dbb825bcbc73dc7401f9846/cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529", size = 523588, upload-time = "2025-09-08T23:24:04.541Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/4b/8d/a0a47a0c9e413a658623d014e91e74a50cdd2c423f7ccfd44086ef767f90/cffi-2.0.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:00bdf7acc5f795150faa6957054fbbca2439db2f775ce831222b66f192f03beb", size = 185230, upload-time = "2025-09-08T23:23:00.879Z" }, + { url = "https://files.pythonhosted.org/packages/4a/d2/a6c0296814556c68ee32009d9c2ad4f85f2707cdecfd7727951ec228005d/cffi-2.0.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:45d5e886156860dc35862657e1494b9bae8dfa63bf56796f2fb56e1679fc0bca", size = 181043, upload-time = "2025-09-08T23:23:02.231Z" }, { url = "https://files.pythonhosted.org/packages/b0/1e/d22cc63332bd59b06481ceaac49d6c507598642e2230f201649058a7e704/cffi-2.0.0-cp313-cp313-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:07b271772c100085dd28b74fa0cd81c8fb1a3ba18b21e03d7c27f3436a10606b", size = 212446, upload-time = "2025-09-08T23:23:03.472Z" }, { url = "https://files.pythonhosted.org/packages/a9/f5/a2c23eb03b61a0b8747f211eb716446c826ad66818ddc7810cc2cc19b3f2/cffi-2.0.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d48a880098c96020b02d5a1f7d9251308510ce8858940e6fa99ece33f610838b", size = 220101, upload-time = "2025-09-08T23:23:04.792Z" }, + { url = "https://files.pythonhosted.org/packages/f2/7f/e6647792fc5850d634695bc0e6ab4111ae88e89981d35ac269956605feba/cffi-2.0.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:f93fd8e5c8c0a4aa1f424d6173f14a892044054871c771f8566e4008eaa359d2", size = 207948, upload-time = "2025-09-08T23:23:06.127Z" }, + { url = "https://files.pythonhosted.org/packages/cb/1e/a5a1bd6f1fb30f22573f76533de12a00bf274abcdc55c8edab639078abb6/cffi-2.0.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:dd4f05f54a52fb558f1ba9f528228066954fee3ebe629fc1660d874d040ae5a3", size = 206422, upload-time = "2025-09-08T23:23:07.753Z" }, { url = "https://files.pythonhosted.org/packages/98/df/0a1755e750013a2081e863e7cd37e0cdd02664372c754e5560099eb7aa44/cffi-2.0.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c8d3b5532fc71b7a77c09192b4a5a200ea992702734a2e9279a37f2478236f26", size = 219499, upload-time = "2025-09-08T23:23:09.648Z" }, { url = "https://files.pythonhosted.org/packages/50/e1/a969e687fcf9ea58e6e2a928ad5e2dd88cc12f6f0ab477e9971f2309b57c/cffi-2.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d9b29c1f0ae438d5ee9acb31cadee00a58c46cc9c0b2f9038c6b0b3470877a8c", size = 222928, upload-time = "2025-09-08T23:23:10.928Z" }, { url = "https://files.pythonhosted.org/packages/36/54/0362578dd2c9e557a28ac77698ed67323ed5b9775ca9d3fe73fe191bb5d8/cffi-2.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6d50360be4546678fc1b79ffe7a66265e28667840010348dd69a314145807a1b", size = 221302, upload-time = "2025-09-08T23:23:12.42Z" }, + { url = "https://files.pythonhosted.org/packages/92/c4/3ce07396253a83250ee98564f8d7e9789fab8e58858f35d07a9a2c78de9f/cffi-2.0.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:fc33c5141b55ed366cfaad382df24fe7dcbc686de5be719b207bb248e3053dc5", size = 185320, upload-time = "2025-09-08T23:23:18.087Z" }, + { url = "https://files.pythonhosted.org/packages/59/dd/27e9fa567a23931c838c6b02d0764611c62290062a6d4e8ff7863daf9730/cffi-2.0.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c654de545946e0db659b3400168c9ad31b5d29593291482c43e3564effbcee13", size = 181487, upload-time = "2025-09-08T23:23:19.622Z" }, { url = "https://files.pythonhosted.org/packages/d6/43/0e822876f87ea8a4ef95442c3d766a06a51fc5298823f884ef87aaad168c/cffi-2.0.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:24b6f81f1983e6df8db3adc38562c83f7d4a0c36162885ec7f7b77c7dcbec97b", size = 220049, upload-time = "2025-09-08T23:23:20.853Z" }, + { url = "https://files.pythonhosted.org/packages/b4/89/76799151d9c2d2d1ead63c2429da9ea9d7aac304603de0c6e8764e6e8e70/cffi-2.0.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:12873ca6cb9b0f0d3a0da705d6086fe911591737a59f28b7936bdfed27c0d47c", size = 207793, upload-time = "2025-09-08T23:23:22.08Z" }, + { url = "https://files.pythonhosted.org/packages/bb/dd/3465b14bb9e24ee24cb88c9e3730f6de63111fffe513492bf8c808a3547e/cffi-2.0.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:d9b97165e8aed9272a6bb17c01e3cc5871a594a446ebedc996e2397a1c1ea8ef", size = 206300, upload-time = "2025-09-08T23:23:23.314Z" }, { url = "https://files.pythonhosted.org/packages/47/d9/d83e293854571c877a92da46fdec39158f8d7e68da75bf73581225d28e90/cffi-2.0.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:afb8db5439b81cf9c9d0c80404b60c3cc9c3add93e114dcae767f1477cb53775", size = 219244, upload-time = "2025-09-08T23:23:24.541Z" }, { url = "https://files.pythonhosted.org/packages/2b/0f/1f177e3683aead2bb00f7679a16451d302c436b5cbf2505f0ea8146ef59e/cffi-2.0.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:737fe7d37e1a1bffe70bd5754ea763a62a066dc5913ca57e957824b72a85e205", size = 222828, upload-time = "2025-09-08T23:23:26.143Z" }, { url = "https://files.pythonhosted.org/packages/c6/0f/cafacebd4b040e3119dcb32fed8bdef8dfe94da653155f9d0b9dc660166e/cffi-2.0.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:38100abb9d1b1435bc4cc340bb4489635dc2f0da7456590877030c9b3d40b0c1", size = 220926, upload-time = "2025-09-08T23:23:27.873Z" }, + { url = "https://files.pythonhosted.org/packages/3e/61/c768e4d548bfa607abcda77423448df8c471f25dbe64fb2ef6d555eae006/cffi-2.0.0-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:9a67fc9e8eb39039280526379fb3a70023d77caec1852002b4da7e8b270c4dd9", size = 188773, upload-time = "2025-09-08T23:23:29.347Z" }, + { url = "https://files.pythonhosted.org/packages/2c/ea/5f76bce7cf6fcd0ab1a1058b5af899bfbef198bea4d5686da88471ea0336/cffi-2.0.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7a66c7204d8869299919db4d5069a82f1561581af12b11b3c9f48c584eb8743d", size = 185013, upload-time = "2025-09-08T23:23:30.63Z" }, { url = "https://files.pythonhosted.org/packages/be/b4/c56878d0d1755cf9caa54ba71e5d049479c52f9e4afc230f06822162ab2f/cffi-2.0.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7cc09976e8b56f8cebd752f7113ad07752461f48a58cbba644139015ac24954c", size = 221593, upload-time = "2025-09-08T23:23:31.91Z" }, + { url = "https://files.pythonhosted.org/packages/e0/0d/eb704606dfe8033e7128df5e90fee946bbcb64a04fcdaa97321309004000/cffi-2.0.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:92b68146a71df78564e4ef48af17551a5ddd142e5190cdf2c5624d0c3ff5b2e8", size = 209354, upload-time = "2025-09-08T23:23:33.214Z" }, + { url = "https://files.pythonhosted.org/packages/d8/19/3c435d727b368ca475fb8742ab97c9cb13a0de600ce86f62eab7fa3eea60/cffi-2.0.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:b1e74d11748e7e98e2f426ab176d4ed720a64412b6a15054378afdb71e0f37dc", size = 208480, upload-time = "2025-09-08T23:23:34.495Z" }, { url = "https://files.pythonhosted.org/packages/d0/44/681604464ed9541673e486521497406fadcc15b5217c3e326b061696899a/cffi-2.0.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:28a3a209b96630bca57cce802da70c266eb08c6e97e5afd61a75611ee6c64592", size = 221584, upload-time = "2025-09-08T23:23:36.096Z" }, { url = "https://files.pythonhosted.org/packages/25/8e/342a504ff018a2825d395d44d63a767dd8ebc927ebda557fecdaca3ac33a/cffi-2.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:7553fb2090d71822f02c629afe6042c299edf91ba1bf94951165613553984512", size = 224443, upload-time = "2025-09-08T23:23:37.328Z" }, { url = "https://files.pythonhosted.org/packages/e1/5e/b666bacbbc60fbf415ba9988324a132c9a7a0448a9a8f125074671c0f2c3/cffi-2.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6c6c373cfc5c83a975506110d17457138c8c63016b563cc9ed6e056a82f13ce4", size = 223437, upload-time = "2025-09-08T23:23:38.945Z" }, @@ -344,8 +356,10 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5c/49/498c86566a1d80e978b42f0d702795f69887005548c041636df6ae1ca64c/cryptography-46.0.3-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:01ca9ff2885f3acc98c29f1860552e37f6d7c7d013d7334ff2a9de43a449315d", size = 4450807, upload-time = "2025-10-15T23:16:56.414Z" }, { url = "https://files.pythonhosted.org/packages/4b/0a/863a3604112174c8624a2ac3c038662d9e59970c7f926acdcfaed8d61142/cryptography-46.0.3-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:6eae65d4c3d33da080cff9c4ab1f711b15c1d9760809dad6ea763f3812d254cb", size = 4299615, upload-time = "2025-10-15T23:16:58.442Z" }, { url = "https://files.pythonhosted.org/packages/64/02/b73a533f6b64a69f3cd3872acb6ebc12aef924d8d103133bb3ea750dc703/cryptography-46.0.3-cp311-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:e5bf0ed4490068a2e72ac03d786693adeb909981cc596425d09032d372bcc849", size = 4016800, upload-time = "2025-10-15T23:17:00.378Z" }, + { url = "https://files.pythonhosted.org/packages/25/d5/16e41afbfa450cde85a3b7ec599bebefaef16b5c6ba4ec49a3532336ed72/cryptography-46.0.3-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:5ecfccd2329e37e9b7112a888e76d9feca2347f12f37918facbb893d7bb88ee8", size = 4984707, upload-time = "2025-10-15T23:17:01.98Z" }, { url = "https://files.pythonhosted.org/packages/c9/56/e7e69b427c3878352c2fb9b450bd0e19ed552753491d39d7d0a2f5226d41/cryptography-46.0.3-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:a2c0cd47381a3229c403062f764160d57d4d175e022c1df84e168c6251a22eec", size = 4482541, upload-time = "2025-10-15T23:17:04.078Z" }, { url = "https://files.pythonhosted.org/packages/78/f6/50736d40d97e8483172f1bb6e698895b92a223dba513b0ca6f06b2365339/cryptography-46.0.3-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:549e234ff32571b1f4076ac269fcce7a808d3bf98b76c8dd560e42dbc66d7d91", size = 4299464, upload-time = "2025-10-15T23:17:05.483Z" }, + { url = "https://files.pythonhosted.org/packages/00/de/d8e26b1a855f19d9994a19c702fa2e93b0456beccbcfe437eda00e0701f2/cryptography-46.0.3-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:c0a7bb1a68a5d3471880e264621346c48665b3bf1c3759d682fc0864c540bd9e", size = 4950838, upload-time = "2025-10-15T23:17:07.425Z" }, { url = "https://files.pythonhosted.org/packages/8f/29/798fc4ec461a1c9e9f735f2fc58741b0daae30688f41b2497dcbc9ed1355/cryptography-46.0.3-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:10b01676fc208c3e6feeb25a8b83d81767e8059e1fe86e1dc62d10a3018fa926", size = 4481596, upload-time = "2025-10-15T23:17:09.343Z" }, { url = "https://files.pythonhosted.org/packages/15/8d/03cd48b20a573adfff7652b76271078e3045b9f49387920e7f1f631d125e/cryptography-46.0.3-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0abf1ffd6e57c67e92af68330d05760b7b7efb243aab8377e583284dbab72c71", size = 4426782, upload-time = "2025-10-15T23:17:11.22Z" }, { url = "https://files.pythonhosted.org/packages/fa/b1/ebacbfe53317d55cf33165bda24c86523497a6881f339f9aae5c2e13e57b/cryptography-46.0.3-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:a04bee9ab6a4da801eb9b51f1b708a1b5b5c9eb48c03f74198464c66f0d344ac", size = 4698381, upload-time = "2025-10-15T23:17:12.829Z" }, @@ -353,8 +367,10 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c5/fd/bc1daf8230eaa075184cbbf5f8cd00ba9db4fd32d63fb83da4671b72ed8a/cryptography-46.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:39b6755623145ad5eff1dab323f4eae2a32a77a7abef2c5089a04a3d04366715", size = 4435078, upload-time = "2025-10-15T23:17:23.042Z" }, { url = "https://files.pythonhosted.org/packages/82/98/d3bd5407ce4c60017f8ff9e63ffee4200ab3e23fe05b765cab805a7db008/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:db391fa7c66df6762ee3f00c95a89e6d428f4d60e7abc8328f4fe155b5ac6e54", size = 4293460, upload-time = "2025-10-15T23:17:24.885Z" }, { url = "https://files.pythonhosted.org/packages/26/e9/e23e7900983c2b8af7a08098db406cf989d7f09caea7897e347598d4cd5b/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:78a97cf6a8839a48c49271cdcbd5cf37ca2c1d6b7fdd86cc864f302b5e9bf459", size = 3995237, upload-time = "2025-10-15T23:17:26.449Z" }, + { url = "https://files.pythonhosted.org/packages/91/15/af68c509d4a138cfe299d0d7ddb14afba15233223ebd933b4bbdbc7155d3/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:dfb781ff7eaa91a6f7fd41776ec37c5853c795d3b358d4896fdbb5df168af422", size = 4967344, upload-time = "2025-10-15T23:17:28.06Z" }, { url = "https://files.pythonhosted.org/packages/ca/e3/8643d077c53868b681af077edf6b3cb58288b5423610f21c62aadcbe99f4/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:6f61efb26e76c45c4a227835ddeae96d83624fb0d29eb5df5b96e14ed1a0afb7", size = 4466564, upload-time = "2025-10-15T23:17:29.665Z" }, { url = "https://files.pythonhosted.org/packages/0e/43/c1e8726fa59c236ff477ff2b5dc071e54b21e5a1e51aa2cee1676f1c986f/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:23b1a8f26e43f47ceb6d6a43115f33a5a37d57df4ea0ca295b780ae8546e8044", size = 4292415, upload-time = "2025-10-15T23:17:31.686Z" }, + { url = "https://files.pythonhosted.org/packages/42/f9/2f8fefdb1aee8a8e3256a0568cffc4e6d517b256a2fe97a029b3f1b9fe7e/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:b419ae593c86b87014b9be7396b385491ad7f320bde96826d0dd174459e54665", size = 4931457, upload-time = "2025-10-15T23:17:33.478Z" }, { url = "https://files.pythonhosted.org/packages/79/30/9b54127a9a778ccd6d27c3da7563e9f2d341826075ceab89ae3b41bf5be2/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:50fc3343ac490c6b08c0cf0d704e881d0d660be923fd3076db3e932007e726e3", size = 4466074, upload-time = "2025-10-15T23:17:35.158Z" }, { url = "https://files.pythonhosted.org/packages/ac/68/b4f4a10928e26c941b1b6a179143af9f4d27d88fe84a6a3c53592d2e76bf/cryptography-46.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:22d7e97932f511d6b0b04f2bfd818d73dcd5928db509460aaf48384778eb6d20", size = 4420569, upload-time = "2025-10-15T23:17:37.188Z" }, { url = "https://files.pythonhosted.org/packages/a3/49/3746dab4c0d1979888f125226357d3262a6dd40e114ac29e3d2abdf1ec55/cryptography-46.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:d55f3dffadd674514ad19451161118fd010988540cee43d8bc20675e775925de", size = 4681941, upload-time = "2025-10-15T23:17:39.236Z" }, @@ -362,13 +378,82 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/26/42/fa8389d4478368743e24e61eea78846a0006caffaf72ea24a15159215a14/cryptography-46.0.3-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:15ab9b093e8f09daab0f2159bb7e47532596075139dd74365da52ecc9cb46c5d", size = 4440029, upload-time = "2025-10-15T23:17:49.837Z" }, { url = "https://files.pythonhosted.org/packages/5f/eb/f483db0ec5ac040824f269e93dd2bd8a21ecd1027e77ad7bdf6914f2fd80/cryptography-46.0.3-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:46acf53b40ea38f9c6c229599a4a13f0d46a6c3fa9ef19fc1a124d62e338dfa0", size = 4297222, upload-time = "2025-10-15T23:17:51.357Z" }, { url = "https://files.pythonhosted.org/packages/fd/cf/da9502c4e1912cb1da3807ea3618a6829bee8207456fbbeebc361ec38ba3/cryptography-46.0.3-cp38-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:10ca84c4668d066a9878890047f03546f3ae0a6b8b39b697457b7757aaf18dbc", size = 4012280, upload-time = "2025-10-15T23:17:52.964Z" }, + { url = "https://files.pythonhosted.org/packages/6b/8f/9adb86b93330e0df8b3dcf03eae67c33ba89958fc2e03862ef1ac2b42465/cryptography-46.0.3-cp38-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:36e627112085bb3b81b19fed209c05ce2a52ee8b15d161b7c643a7d5a88491f3", size = 4978958, upload-time = "2025-10-15T23:17:54.965Z" }, { url = "https://files.pythonhosted.org/packages/d1/a0/5fa77988289c34bdb9f913f5606ecc9ada1adb5ae870bd0d1054a7021cc4/cryptography-46.0.3-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:1000713389b75c449a6e979ffc7dcc8ac90b437048766cef052d4d30b8220971", size = 4473714, upload-time = "2025-10-15T23:17:56.754Z" }, { url = "https://files.pythonhosted.org/packages/14/e5/fc82d72a58d41c393697aa18c9abe5ae1214ff6f2a5c18ac470f92777895/cryptography-46.0.3-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:b02cf04496f6576afffef5ddd04a0cb7d49cf6be16a9059d793a30b035f6b6ac", size = 4296970, upload-time = "2025-10-15T23:17:58.588Z" }, + { url = "https://files.pythonhosted.org/packages/78/06/5663ed35438d0b09056973994f1aec467492b33bd31da36e468b01ec1097/cryptography-46.0.3-cp38-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:71e842ec9bc7abf543b47cf86b9a743baa95f4677d22baa4c7d5c69e49e9bc04", size = 4940236, upload-time = "2025-10-15T23:18:00.897Z" }, { url = "https://files.pythonhosted.org/packages/fc/59/873633f3f2dcd8a053b8dd1d38f783043b5fce589c0f6988bf55ef57e43e/cryptography-46.0.3-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:402b58fc32614f00980b66d6e56a5b4118e6cb362ae8f3fda141ba4689bd4506", size = 4472642, upload-time = "2025-10-15T23:18:02.749Z" }, { url = "https://files.pythonhosted.org/packages/3d/39/8e71f3930e40f6877737d6f69248cf74d4e34b886a3967d32f919cc50d3b/cryptography-46.0.3-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ef639cb3372f69ec44915fafcd6698b6cc78fbe0c2ea41be867f6ed612811963", size = 4423126, upload-time = "2025-10-15T23:18:04.85Z" }, { url = "https://files.pythonhosted.org/packages/cd/c7/f65027c2810e14c3e7268353b1681932b87e5a48e65505d8cc17c99e36ae/cryptography-46.0.3-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:3b51b8ca4f1c6453d8829e1eb7299499ca7f313900dd4d89a24b8b87c0a780d4", size = 4686573, upload-time = "2025-10-15T23:18:06.908Z" }, ] +[[package]] +name = "cuda-bindings" +version = "13.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder", marker = "sys_platform == 'linux'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/df/93/eef988860a3ca985f82c4f3174fc0cdd94e07331ba9a92e8e064c260337f/cuda_bindings-13.2.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6629ca2df6f795b784752409bcaedbd22a7a651b74b56a165ebc0c9dcbd504d0", size = 5614610, upload-time = "2026-03-11T00:12:50.337Z" }, + { url = "https://files.pythonhosted.org/packages/18/23/6db3aba46864aee357ab2415135b3fe3da7e9f1fa0221fa2a86a5968099c/cuda_bindings-13.2.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7dca0da053d3b4cc4869eff49c61c03f3c5dbaa0bcd712317a358d5b8f3f385d", size = 6149914, upload-time = "2026-03-11T00:12:52.374Z" }, + { url = "https://files.pythonhosted.org/packages/c0/87/87a014f045b77c6de5c8527b0757fe644417b184e5367db977236a141602/cuda_bindings-13.2.0-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a6464b30f46692d6c7f65d4a0e0450d81dd29de3afc1bb515653973d01c2cd6e", size = 5685673, upload-time = "2026-03-11T00:12:56.371Z" }, + { url = "https://files.pythonhosted.org/packages/ee/5e/c0fe77a73aaefd3fff25ffaccaac69c5a63eafdf8b9a4c476626ef0ac703/cuda_bindings-13.2.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f4af9f3e1be603fa12d5ad6cfca7844c9d230befa9792b5abdf7dd79979c3626", size = 6191386, upload-time = "2026-03-11T00:12:58.965Z" }, + { url = "https://files.pythonhosted.org/packages/5f/58/ed2c3b39c8dd5f96aa7a4abef0d47a73932c7a988e30f5fa428f00ed0da1/cuda_bindings-13.2.0-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:df850a1ff8ce1b3385257b08e47b70e959932f5f432d0a4e46a355962b4e4771", size = 5507469, upload-time = "2026-03-11T00:13:04.063Z" }, + { url = "https://files.pythonhosted.org/packages/1f/01/0c941b112ceeb21439b05895eace78ca1aa2eaaf695c8521a068fd9b4c00/cuda_bindings-13.2.0-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e8a16384c6494e5485f39314b0b4afb04bee48d49edb16d5d8593fd35bbd231b", size = 6059693, upload-time = "2026-03-11T00:13:06.003Z" }, +] + +[[package]] +name = "cuda-pathfinder" +version = "1.5.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/93/66/0c02bd330e7d976f83fa68583d6198d76f23581bcbb5c0e98a6148f326e5/cuda_pathfinder-1.5.0-py3-none-any.whl", hash = "sha256:498f90a9e9de36044a7924742aecce11c50c49f735f1bc53e05aa46de9ea4110", size = 49739, upload-time = "2026-03-24T21:14:30.869Z" }, +] + +[[package]] +name = "cuda-toolkit" +version = "13.0.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/57/b2/453099f5f3b698d7d0eab38916aac44c7f76229f451709e2eb9db6615dcd/cuda_toolkit-13.0.2-py2.py3-none-any.whl", hash = "sha256:b198824cf2f54003f50d64ada3a0f184b42ca0846c1c94192fa269ecd97a66eb", size = 2364, upload-time = "2025-12-19T23:24:07.328Z" }, +] + +[package.optional-dependencies] +cublas = [ + { name = "nvidia-cublas", marker = "sys_platform == 'linux'" }, +] +cudart = [ + { name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux'" }, +] +cufft = [ + { name = "nvidia-cufft", marker = "sys_platform == 'linux'" }, +] +cufile = [ + { name = "nvidia-cufile", marker = "sys_platform == 'linux'" }, +] +cupti = [ + { name = "nvidia-cuda-cupti", marker = "sys_platform == 'linux'" }, +] +curand = [ + { name = "nvidia-curand", marker = "sys_platform == 'linux'" }, +] +cusolver = [ + { name = "nvidia-cusolver", marker = "sys_platform == 'linux'" }, +] +cusparse = [ + { name = "nvidia-cusparse", marker = "sys_platform == 'linux'" }, +] +nvjitlink = [ + { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux'" }, +] +nvrtc = [ + { name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'linux'" }, +] +nvtx = [ + { name = "nvidia-nvtx", marker = "sys_platform == 'linux'" }, +] + [[package]] name = "cycler" version = "0.12.1" @@ -473,8 +558,9 @@ dependencies = [ { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cpu"], marker = "sys_platform == 'linux'" }, - { name = "mlx", version = "0.31.2.dev20260327+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" }, + { name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" }, { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "mlx-vlm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "openai-harmony", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "psutil", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, @@ -483,6 +569,7 @@ dependencies = [ { name = "rustworkx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "tiktoken", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "tomlkit", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "transformers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "types-aiofiles", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "zstandard", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, ] @@ -506,13 +593,14 @@ requires-dist = [ { name = "fastapi", specifier = ">=0.116.1" }, { name = "filelock", specifier = ">=3.18.0" }, { name = "httpx", specifier = ">=0.28.1" }, - { name = "huggingface-hub", specifier = ">=0.33.4" }, + { name = "huggingface-hub", specifier = ">=1.8.0" }, { name = "hypercorn", specifier = ">=0.18.0" }, { name = "loguru", specifier = ">=0.7.3" }, { name = "mflux", specifier = "==0.17.2" }, { name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" }, { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.6" }, { name = "mlx-lm", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-deepseek-v32-indexer" }, + { name = "mlx-vlm", specifier = ">=0.3.11" }, { name = "msgspec", specifier = ">=0.19.0" }, { name = "openai-harmony", specifier = ">=0.0.8" }, { name = "psutil", specifier = ">=7.0.0" }, @@ -521,6 +609,7 @@ requires-dist = [ { name = "rustworkx", specifier = ">=0.17.1" }, { name = "tiktoken", specifier = ">=0.12.0" }, { name = "tomlkit", specifier = ">=0.14.0" }, + { name = "transformers", specifier = ">=5.0.0,<5.4.0" }, { name = "types-aiofiles", specifier = ">=24.1.0.20250708" }, { name = "zstandard", specifier = ">=0.23.0" }, ] @@ -783,31 +872,28 @@ wheels = [ [[package]] name = "hf-xet" -version = "1.2.1rc0" +version = "1.4.2" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/9a/48/61907d37a180a1d016cb79396215b1064f075965cf14ac78b4a9682705d7/hf_xet-1.2.1rc0.tar.gz", hash = "sha256:ee6b196855720767283dbbca6d5f3877afdfa6df83e037bbadbed0181ac5972e", size = 518988, upload-time = "2025-11-21T23:26:10.526Z" } +sdist = { url = "https://files.pythonhosted.org/packages/09/08/23c84a26716382c89151b5b447b4beb19e3345f3a93d3b73009a71a57ad3/hf_xet-1.4.2.tar.gz", hash = "sha256:b7457b6b482d9e0743bd116363239b1fa904a5e65deede350fbc0c4ea67c71ea", size = 672357, upload-time = "2026-03-13T06:58:51.077Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/8c/2b/e9fb76e7dcba1efc0dc881124d0ebbdf0790ad78f90dae9f23a969224c0c/hf_xet-1.2.1rc0-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:05acfd78c5b515a0c06103c9471208a71ae52c6a72dba73bbcb5b7f79575c530", size = 2973766, upload-time = "2025-11-21T23:25:50.546Z" }, - { url = "https://files.pythonhosted.org/packages/95/bf/8365816fb0e2dc0db633bed504fdf70b4e4e052aa86caff62e4b0175e7fa/hf_xet-1.2.1rc0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:2e4bbe0e4195c48aebce7c87438df6ba0748001c15cd088d1f41553b9cbf0aa5", size = 2850724, upload-time = "2025-11-21T23:25:48.95Z" }, - { url = "https://files.pythonhosted.org/packages/4a/52/72ba543089817fdf0e684032c1664fd249602896d52b76f4278b7c830cc8/hf_xet-1.2.1rc0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:66534e7190bafae92c8e3411011220f189fadcc8cba36ebf4bc261e769fb7e49", size = 3342204, upload-time = "2025-11-21T23:25:31.773Z" }, - { url = "https://files.pythonhosted.org/packages/85/a0/d0f7b4ffb08bdb25db2dbad8e5d97a266a4ada3c7e8dc4429bfe99c86ed2/hf_xet-1.2.1rc0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c9d193015364fb9e95d4d295722538b554e9bfaa7b6a167e09e030148c8b15d0", size = 19434060, upload-time = "2025-11-21T23:25:33.89Z" }, - { url = "https://files.pythonhosted.org/packages/af/b4/c406e62a1895520da504bb9372f7ed26ef65e32e1b39e397d81b7136b5ab/hf_xet-1.2.1rc0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:dda4a029cd30f10ba205d8a74e232070ec75923e4c262a2d7f5d55eb3a3dd4d1", size = 3249296, upload-time = "2025-11-21T23:25:29.504Z" }, - { url = "https://files.pythonhosted.org/packages/cf/fb/c40487744c12a038e31af75de661938a6e9c2cfb55a544706d9b9d3cc00c/hf_xet-1.2.1rc0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:fc95e2b7a1a3a613587f407a8292f1240d45febd66a49ee1da0a94414ff3784e", size = 3434401, upload-time = "2025-11-21T23:25:59.747Z" }, - { url = "https://files.pythonhosted.org/packages/46/37/8b93e82bace53bb650474562487a4fe2aa43e8b8d9ecd01ddffc1b6a63f2/hf_xet-1.2.1rc0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:4a4e981ef129bdf1af7be559319b017bed0ae997c8bdd696b6c7e50d888e5a51", size = 3520042, upload-time = "2025-11-21T23:26:01.691Z" }, - { url = "https://files.pythonhosted.org/packages/c7/b7/6ce9f48be8748b2e8599453dec7012d38e4685a5e5587ee3ef4c09fccaf9/hf_xet-1.2.1rc0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:1d57ee9323fcf87c3fc1840856ad2f767c0f8ee14a55d470ddba3a6fdab40dd2", size = 2973781, upload-time = "2025-11-21T23:25:58.073Z" }, - { url = "https://files.pythonhosted.org/packages/72/dc/6e1d3b653fdb34ce86f7b94c2388270f8bb5bb18da8590425a30ef0af1be/hf_xet-1.2.1rc0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:6163f7de633ac0f5f88dc24d369b30df4df0f923dc61ebd9c39a9b022497f47f", size = 2850462, upload-time = "2025-11-21T23:25:56.157Z" }, - { url = "https://files.pythonhosted.org/packages/8c/6b/6e0daf5811badf6c9d60a49cb3f99fe41cc01f147ecae3911d8621fa69c1/hf_xet-1.2.1rc0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:05b518a2499dafd510e29ff6c14bfb9aae119f66af785fc99eaf9069e0ccda43", size = 3342036, upload-time = "2025-11-21T23:25:44.283Z" }, - { url = "https://files.pythonhosted.org/packages/b7/21/9dfdf0c66743cbf14f312d196f19367372a89232b2623d733690474008b9/hf_xet-1.2.1rc0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e5ee726b80a1c0b2868bc58302ba1a47d0702f8d67f69aeecb94fe7f30ac1c2b", size = 19431002, upload-time = "2025-11-21T23:25:46.621Z" }, - { url = "https://files.pythonhosted.org/packages/f4/8c/f798608de78b5aa1cabbf9c1e5e8a0172a93a47267fe1733f7c9780802e2/hf_xet-1.2.1rc0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:bf8f5439c39a5fa41dec1071f9576ac510180522690771d54c211151e08cdf35", size = 3248725, upload-time = "2025-11-21T23:25:42.387Z" }, - { url = "https://files.pythonhosted.org/packages/75/75/7035ea757b2ef27c21a7d734da18c1537473f8dcff468872eb9b4281dd33/hf_xet-1.2.1rc0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:5ca1fae9189095b15c89cd30ce2f6c3a97f2d1cab261e28a73b84690ebc8960a", size = 3433685, upload-time = "2025-11-21T23:26:06.88Z" }, - { url = "https://files.pythonhosted.org/packages/0e/47/1627f85cb062283edc9f516d61838c88bcdb46828d903b035674b5e0e89c/hf_xet-1.2.1rc0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:99676d52bbffc7747950d2686bc91f520758f3d83b594988058478be68706862", size = 3519636, upload-time = "2025-11-21T23:26:08.512Z" }, - { url = "https://files.pythonhosted.org/packages/6e/ce/bfd825a3aa2a22caa78865a6331e3660825b82de24877b08c10d18c45748/hf_xet-1.2.1rc0-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:b6b6455d68f2b4439028c58198e6dc33f3b1b64314ed05b0a5f5f7dace37d711", size = 2977924, upload-time = "2025-11-21T23:25:54.254Z" }, - { url = "https://files.pythonhosted.org/packages/88/28/d78d7fcf2f3e18177e8dd6bbb4294bb00ef2f6d3addfc2b636a251ec297b/hf_xet-1.2.1rc0-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:3d9894128c63478a3f67d7f0288e8f5780c2b3ae7a09f36fc3949be60dcf7ac8", size = 2853755, upload-time = "2025-11-21T23:25:52.222Z" }, - { url = "https://files.pythonhosted.org/packages/ae/09/637245509430b3dd9d37f676bbe0b993c723e3671ce0b39fdf42c6f05a02/hf_xet-1.2.1rc0-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f8b937c5e2a4f43720eca9564b14324ecfa108cc053a1b44890c620f51aac01e", size = 3347297, upload-time = "2025-11-21T23:25:37.9Z" }, - { url = "https://files.pythonhosted.org/packages/29/b5/bbc98a35ee5229d0cd6c9436ae97f86cf2ab63d6bd463cd5a43282e5c1f8/hf_xet-1.2.1rc0-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7bd4629e923dd7b12fb9d05312e03ed123db230ae25fd98a3fd5caa739f2357e", size = 19457253, upload-time = "2025-11-21T23:25:40.115Z" }, - { url = "https://files.pythonhosted.org/packages/0f/c6/ab21fc91f23ca54cdd44e86981d80475d67ee4122128f5ef988a119ebe28/hf_xet-1.2.1rc0-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:5484ad943ceec043f0c29733cb87e59c86c2c68804c470176f259b1ef339718e", size = 3254771, upload-time = "2025-11-21T23:25:36.213Z" }, - { url = "https://files.pythonhosted.org/packages/e6/c0/5a2887739722bd5a531769c1e9555e30dd7f470aefaabbe898d939dbba20/hf_xet-1.2.1rc0-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:2ec943ba2633ed0df48d2c817ce6a13670e96590f9fd4260011c5753afbc5d53", size = 3439600, upload-time = "2025-11-21T23:26:03.318Z" }, - { url = "https://files.pythonhosted.org/packages/30/c9/c7cd0a64eb2dba1f70fbb78dee33558567404522776328254a7c805ae23e/hf_xet-1.2.1rc0-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:87e0bdd71172b7cb1621e706bbf70b75f31df5fa7c359ebc0978567b5c21c2cf", size = 3526094, upload-time = "2025-11-21T23:26:05.018Z" }, + { url = "https://files.pythonhosted.org/packages/18/06/e8cf74c3c48e5485c7acc5a990d0d8516cdfb5fdf80f799174f1287cc1b5/hf_xet-1.4.2-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:ac8202ae1e664b2c15cdfc7298cbb25e80301ae596d602ef7870099a126fcad4", size = 3796125, upload-time = "2026-03-13T06:58:33.177Z" }, + { url = "https://files.pythonhosted.org/packages/66/d4/b73ebab01cbf60777323b7de9ef05550790451eb5172a220d6b9845385ec/hf_xet-1.4.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:6d2f8ee39fa9fba9af929f8c0d0482f8ee6e209179ad14a909b6ad78ffcb7c81", size = 3555985, upload-time = "2026-03-13T06:58:31.797Z" }, + { url = "https://files.pythonhosted.org/packages/ff/e7/ded6d1bd041c3f2bca9e913a0091adfe32371988e047dd3a68a2463c15a2/hf_xet-1.4.2-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4642a6cf249c09da8c1f87fe50b24b2a3450b235bf8adb55700b52f0ea6e2eb6", size = 4212085, upload-time = "2026-03-13T06:58:24.323Z" }, + { url = "https://files.pythonhosted.org/packages/97/c1/a0a44d1f98934f7bdf17f7a915b934f9fca44bb826628c553589900f6df8/hf_xet-1.4.2-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:769431385e746c92dc05492dde6f687d304584b89c33d79def8367ace06cb555", size = 3988266, upload-time = "2026-03-13T06:58:22.887Z" }, + { url = "https://files.pythonhosted.org/packages/7a/82/be713b439060e7d1f1d93543c8053d4ef2fe7e6922c5b31642eaa26f3c4b/hf_xet-1.4.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:c9dd1c1bc4cc56168f81939b0e05b4c36dd2d28c13dc1364b17af89aa0082496", size = 4188513, upload-time = "2026-03-13T06:58:40.858Z" }, + { url = "https://files.pythonhosted.org/packages/21/a6/cbd4188b22abd80ebd0edbb2b3e87f2633e958983519980815fb8314eae5/hf_xet-1.4.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:fca58a2ae4e6f6755cc971ac6fcdf777ea9284d7e540e350bb000813b9a3008d", size = 4428287, upload-time = "2026-03-13T06:58:42.601Z" }, + { url = "https://files.pythonhosted.org/packages/1e/0f/fcd2504015eab26358d8f0f232a1aed6b8d363a011adef83fe130bff88f7/hf_xet-1.4.2-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:949dcf88b484bb9d9276ca83f6599e4aa03d493c08fc168c124ad10b2e6f75d7", size = 3796493, upload-time = "2026-03-13T06:58:39.267Z" }, + { url = "https://files.pythonhosted.org/packages/82/56/19c25105ff81731ca6d55a188b5de2aa99d7a2644c7aa9de1810d5d3b726/hf_xet-1.4.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:41659966020d59eb9559c57de2cde8128b706a26a64c60f0531fa2318f409418", size = 3555797, upload-time = "2026-03-13T06:58:37.546Z" }, + { url = "https://files.pythonhosted.org/packages/bf/e3/8933c073186849b5e06762aa89847991d913d10a95d1603eb7f2c3834086/hf_xet-1.4.2-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5c588e21d80010119458dd5d02a69093f0d115d84e3467efe71ffb2c67c19146", size = 4212127, upload-time = "2026-03-13T06:58:30.539Z" }, + { url = "https://files.pythonhosted.org/packages/eb/01/f89ebba4e369b4ed699dcb60d3152753870996f41c6d22d3d7cac01310e1/hf_xet-1.4.2-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:a296744d771a8621ad1d50c098d7ab975d599800dae6d48528ba3944e5001ba0", size = 3987788, upload-time = "2026-03-13T06:58:29.139Z" }, + { url = "https://files.pythonhosted.org/packages/84/4d/8a53e5ffbc2cc33bbf755382ac1552c6d9af13f623ed125fe67cc3e6772f/hf_xet-1.4.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:f563f7efe49588b7d0629d18d36f46d1658fe7e08dce3fa3d6526e1c98315e2d", size = 4188315, upload-time = "2026-03-13T06:58:48.017Z" }, + { url = "https://files.pythonhosted.org/packages/d1/b8/b7a1c1b5592254bd67050632ebbc1b42cc48588bf4757cb03c2ef87e704a/hf_xet-1.4.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5b2e0132c56d7ee1bf55bdb638c4b62e7106f6ac74f0b786fed499d5548c5570", size = 4428306, upload-time = "2026-03-13T06:58:49.502Z" }, + { url = "https://files.pythonhosted.org/packages/b4/86/b40b83a2ff03ef05c4478d2672b1fc2b9683ff870e2b25f4f3af240f2e7b/hf_xet-1.4.2-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:71f02d6e4cdd07f344f6844845d78518cc7186bd2bc52d37c3b73dc26a3b0bc5", size = 3800339, upload-time = "2026-03-13T06:58:36.245Z" }, + { url = "https://files.pythonhosted.org/packages/64/2e/af4475c32b4378b0e92a587adb1aa3ec53e3450fd3e5fe0372a874531c00/hf_xet-1.4.2-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:e9b38d876e94d4bdcf650778d6ebbaa791dd28de08db9736c43faff06ede1b5a", size = 3559664, upload-time = "2026-03-13T06:58:34.787Z" }, + { url = "https://files.pythonhosted.org/packages/3c/4c/781267da3188db679e601de18112021a5cb16506fe86b246e22c5401a9c4/hf_xet-1.4.2-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:77e8c180b7ef12d8a96739a4e1e558847002afe9ea63b6f6358b2271a8bdda1c", size = 4217422, upload-time = "2026-03-13T06:58:27.472Z" }, + { url = "https://files.pythonhosted.org/packages/68/47/d6cf4a39ecf6c7705f887a46f6ef5c8455b44ad9eb0d391aa7e8a2ff7fea/hf_xet-1.4.2-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:c3b3c6a882016b94b6c210957502ff7877802d0dbda8ad142c8595db8b944271", size = 3992847, upload-time = "2026-03-13T06:58:25.989Z" }, + { url = "https://files.pythonhosted.org/packages/2d/ef/e80815061abff54697239803948abc665c6b1d237102c174f4f7a9a5ffc5/hf_xet-1.4.2-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:9d9a634cc929cfbaf2e1a50c0e532ae8c78fa98618426769480c58501e8c8ac2", size = 4193843, upload-time = "2026-03-13T06:58:44.59Z" }, + { url = "https://files.pythonhosted.org/packages/54/75/07f6aa680575d9646c4167db6407c41340cbe2357f5654c4e72a1b01ca14/hf_xet-1.4.2-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:6b0932eb8b10317ea78b7da6bab172b17be03bbcd7809383d8d5abd6a2233e04", size = 4432751, upload-time = "2026-03-13T06:58:46.533Z" }, ] [[package]] @@ -849,7 +935,7 @@ wheels = [ [[package]] name = "huggingface-hub" -version = "1.3.1" +version = "1.8.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, @@ -858,14 +944,13 @@ dependencies = [ { name = "httpx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "packaging", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "shellingham", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "tqdm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "typer-slim", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "typer", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/dd/dd/1cc985c5dda36298b152f75e82a1c81f52243b78fb7e9cad637a29561ad1/huggingface_hub-1.3.1.tar.gz", hash = "sha256:e80e0cfb4a75557c51ab20d575bdea6bb6106c2f97b7c75d8490642f1efb6df5", size = 622356, upload-time = "2026-01-09T14:08:16.888Z" } +sdist = { url = "https://files.pythonhosted.org/packages/8e/2a/a847fd02261cd051da218baf99f90ee7c7040c109a01833db4f838f25256/huggingface_hub-1.8.0.tar.gz", hash = "sha256:c5627b2fd521e00caf8eff4ac965ba988ea75167fad7ee72e17f9b7183ec63f3", size = 735839, upload-time = "2026-03-25T16:01:28.152Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/90/fb/cb8fe5f71d5622427f20bcab9e06a696a5aaf21bfe7bd0a8a0c63c88abf5/huggingface_hub-1.3.1-py3-none-any.whl", hash = "sha256:efbc7f3153cb84e2bb69b62ed90985e21ecc9343d15647a419fc0ee4b85f0ac3", size = 533351, upload-time = "2026-01-09T14:08:14.519Z" }, + { url = "https://files.pythonhosted.org/packages/a9/ae/8a3a16ea4d202cb641b51d2681bdd3d482c1c592d7570b3fa264730829ce/huggingface_hub-1.8.0-py3-none-any.whl", hash = "sha256:d3eb5047bd4e33c987429de6020d4810d38a5bef95b3b40df9b17346b7f353f2", size = 625208, upload-time = "2026-03-25T16:01:26.603Z" }, ] [[package]] @@ -1351,7 +1436,7 @@ dependencies = [ { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform == 'linux'" }, - { name = "mlx", version = "0.31.2.dev20260327+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" }, + { name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" }, { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "piexif", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, @@ -1375,6 +1460,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/68/02/f94eca4e77b7d12685060461eb793cbc8c00e96cc7fe0ce376374201aed2/mflux-0.17.2-py3-none-any.whl", hash = "sha256:be1642b04847413c0a8ed1dae82ce1ca023e155b057d82a8301eca9c3fe08339", size = 1037451, upload-time = "2026-03-23T13:08:16.747Z" }, ] +[[package]] +name = "miniaudio" +version = "1.61" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cffi", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/55/fa/96d4cc7ada283357117f7890418ac065a0a6d81ec59e681cd965a403aba3/miniaudio-1.61.tar.gz", hash = "sha256:e88e97837d031f0fb6982394218b6487de02eaa382ad273b8fca37791a2b4b15", size = 1103527, upload-time = "2024-07-24T18:13:10.037Z" } + [[package]] name = "mlx" version = "0.30.6" @@ -1400,7 +1494,7 @@ cuda13 = [ [[package]] name = "mlx" -version = "0.31.2.dev20260327+e5e64331" +version = "0.31.2.dev20260324+e5e64331" source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" } resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'darwin'", @@ -1437,7 +1531,7 @@ version = "0.31.2" source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-deepseek-v32-indexer#d388ff77858fec3b5d2e3b1d9502a7e2878b8109" } dependencies = [ { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "mlx", version = "0.31.2.dev20260327+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" }, + { name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" }, { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, @@ -1445,6 +1539,30 @@ dependencies = [ { name = "transformers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, ] +[[package]] +name = "mlx-vlm" +version = "0.4.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "datasets", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "miniaudio", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux'" }, + { name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" }, + { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "pillow", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "requests", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "tqdm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "transformers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "uvicorn", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/57/8f/31204f1a8c7404e523a5578949ea668e668e10dd67a0f63336f261014c0f/mlx_vlm-0.4.1.tar.gz", hash = "sha256:4e2d8a232715dbca72d346f43cf54d5738452848855792ffb1b285228ae7c7bd", size = 621840, upload-time = "2026-03-21T14:26:04.586Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/78/856f44c6bdd8791427fa59a093a1d00c91cdbe16238506602fc3968017bb/mlx_vlm-0.4.1-py3-none-any.whl", hash = "sha256:89feca2e8be31609770c0e8a6d88fa21d00ee25bd3d56b4aafce59d35dd63b71", size = 768806, upload-time = "2026-03-21T14:26:03.129Z" }, +] + [[package]] name = "more-itertools" version = "10.8.0" @@ -1680,175 +1798,151 @@ wheels = [ [[package]] name = "nvidia-cublas" -version = "13.2.1.1" +version = "13.1.0.3" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9d/36/0124129e1378e9834e0cbe19781fbe0ffd5f870c2af6f01cdf17a9869c39/nvidia_cublas-13.2.1.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:8b4a4cd8b73772fde9ccaa1f3967eb001ae5fde8b1dc37f7442d072b64d6f5da", size = 502470979, upload-time = "2026-01-13T22:39:37.619Z" }, - { url = "https://files.pythonhosted.org/packages/e2/e7/39e43c0688f9788c88da0b91ea18125448c5f515104aadf65a70243f144f/nvidia_cublas-13.2.1.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8c13c93cf8be4480b4909905c96d2d31575b4af43fcd3af0e84af94762665e4f", size = 401085577, upload-time = "2026-01-13T22:40:18.702Z" }, + { url = "https://files.pythonhosted.org/packages/e1/a5/fce49e2ae977e0ccc084e5adafceb4f0ac0c8333cb6863501618a7277f67/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:c86fc7f7ae36d7528288c5d88098edcb7b02c633d262e7ddbb86b0ad91be5df2", size = 542851226, upload-time = "2025-10-09T08:59:04.818Z" }, + { url = "https://files.pythonhosted.org/packages/e7/44/423ac00af4dd95a5aeb27207e2c0d9b7118702149bf4704c3ddb55bb7429/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:ee8722c1f0145ab246bccb9e452153b5e0515fd094c3678df50b2a0888b8b171", size = 423133236, upload-time = "2025-10-09T08:59:32.536Z" }, ] [[package]] -name = "nvidia-cublas-cu12" -version = "12.8.4.1" +name = "nvidia-cuda-cupti" +version = "13.0.85" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8ac4e771d5a348c551b2a426eda6193c19aa630236b418086020df5ba9667142", size = 594346921, upload-time = "2025-03-07T01:44:31.254Z" }, -] - -[[package]] -name = "nvidia-cuda-cupti-cu12" -version = "12.8.90" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f8/02/2adcaa145158bf1a8295d83591d22e4103dbfd821bcaf6f3f53151ca4ffa/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ea0cb07ebda26bb9b29ba82cda34849e73c166c18162d3913575b0c9db9a6182", size = 10248621, upload-time = "2025-03-07T01:40:21.213Z" }, + { url = "https://files.pythonhosted.org/packages/2a/2a/80353b103fc20ce05ef51e928daed4b6015db4aaa9162ed0997090fe2250/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:796bd679890ee55fb14a94629b698b6db54bcfd833d391d5e94017dd9d7d3151", size = 10310827, upload-time = "2025-09-04T08:26:42.012Z" }, + { url = "https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8", size = 10715597, upload-time = "2025-09-04T08:26:51.312Z" }, ] [[package]] name = "nvidia-cuda-nvrtc" -version = "13.1.115" +version = "13.0.88" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ff/d8/6fcf0f32d133a7da92efb1e90844d9f7c104627066cc52b13f7f0b128b54/nvidia_cuda_nvrtc-13.1.115-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:d7cf1284ab82f379884decc8813d9d4a729bc96b1ec020a9cf80303f698d73c4", size = 46564545, upload-time = "2026-01-13T22:35:53.834Z" }, - { url = "https://files.pythonhosted.org/packages/e9/f8/8bab039cbdd87af53f2ca0ca9e93bd676e53393ab4ea43da4735854dc1ce/nvidia_cuda_nvrtc-13.1.115-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:dfbc5e3bb19db41e4a05280b7b0cb9cbb624699f57dab3798455f43345541f99", size = 44308134, upload-time = "2026-01-13T22:35:35.287Z" }, + { url = "https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575", size = 90215200, upload-time = "2025-09-04T08:28:44.204Z" }, + { url = "https://files.pythonhosted.org/packages/b7/dc/6bb80850e0b7edd6588d560758f17e0550893a1feaf436807d64d2da040f/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b", size = 43015449, upload-time = "2025-09-04T08:28:20.239Z" }, ] [[package]] -name = "nvidia-cuda-nvrtc-cu12" -version = "12.8.93" +name = "nvidia-cuda-runtime" +version = "13.0.96" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/05/6b/32f747947df2da6994e999492ab306a903659555dddc0fbdeb9d71f75e52/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994", size = 88040029, upload-time = "2025-03-07T01:42:13.562Z" }, -] - -[[package]] -name = "nvidia-cuda-runtime-cu12" -version = "12.8.90" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/0d/9b/a997b638fcd068ad6e4d53b8551a7d30fe8b404d6f1804abf1df69838932/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90", size = 954765, upload-time = "2025-03-07T01:40:01.615Z" }, -] - -[[package]] -name = "nvidia-cudnn-cu12" -version = "9.10.2.21" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" }, + { url = "https://files.pythonhosted.org/packages/87/4f/17d7b9b8e285199c58ce28e31b5c5bbaa4d8271af06a89b6405258245de2/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55", size = 2261060, upload-time = "2025-10-09T08:55:15.78Z" }, + { url = "https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548", size = 2243632, upload-time = "2025-10-09T08:55:36.117Z" }, ] [[package]] name = "nvidia-cudnn-cu13" -version = "9.18.0.77" +version = "9.19.0.56" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-cublas", marker = "sys_platform == 'linux'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/0b/74/ad8b8edef8b8a54071cb8bd80b63aee7a833b1eabfcdbc0bbec4f0868cc1/nvidia_cudnn_cu13-9.18.0.77-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:a0957ae96e752840033a505bb6505d634e03c4bb4947e3e8fe1fdbe599120ab3", size = 423102639, upload-time = "2026-01-16T20:25:15.997Z" }, - { url = "https://files.pythonhosted.org/packages/d7/e9/c26231f84cc906dca63f4517e9824c9e8d166838e61fc0c5015cbe11fe59/nvidia_cudnn_cu13-9.18.0.77-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:82b86573bd24bfc1450c2e94866af23f5327982a7a4f14ee2416f8e0dd4631f6", size = 358843113, upload-time = "2026-01-16T20:26:35.975Z" }, + { url = "https://files.pythonhosted.org/packages/f1/84/26025437c1e6b61a707442184fa0c03d083b661adf3a3eecfd6d21677740/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:6ed29ffaee1176c612daf442e4dd6cfeb6a0caa43ddcbeb59da94953030b1be4", size = 433781201, upload-time = "2026-02-03T20:40:53.805Z" }, + { url = "https://files.pythonhosted.org/packages/a3/22/0b4b932655d17a6da1b92fa92ab12844b053bb2ac2475e179ba6f043da1e/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:d20e1734305e9d68889a96e3f35094d733ff1f83932ebe462753973e53a572bf", size = 366066321, upload-time = "2026-02-03T20:44:52.837Z" }, ] [[package]] -name = "nvidia-cufft-cu12" -version = "11.3.3.83" +name = "nvidia-cufft" +version = "12.0.0.61" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" }, + { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, + { url = "https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3", size = 214085489, upload-time = "2025-09-04T08:31:56.044Z" }, ] [[package]] -name = "nvidia-cufile-cu12" -version = "1.13.1.3" +name = "nvidia-cufile" +version = "1.15.1.6" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/bb/fe/1bcba1dfbfb8d01be8d93f07bfc502c93fa23afa6fd5ab3fc7c1df71038a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d069003be650e131b21c932ec3d8969c1715379251f8d23a1860554b1cb24fc", size = 1197834, upload-time = "2025-03-07T01:45:50.723Z" }, + { url = "https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44", size = 1223672, upload-time = "2025-09-04T08:32:22.779Z" }, + { url = "https://files.pythonhosted.org/packages/ab/73/cc4a14c9813a8a0d509417cf5f4bdaba76e924d58beb9864f5a7baceefbf/nvidia_cufile-1.15.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:bdc0deedc61f548bddf7733bdc216456c2fdb101d020e1ab4b88d232d5e2f6d1", size = 1136992, upload-time = "2025-09-04T08:32:14.119Z" }, ] [[package]] -name = "nvidia-curand-cu12" -version = "10.3.9.90" +name = "nvidia-curand" +version = "10.4.0.35" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/fb/aa/6584b56dc84ebe9cf93226a5cde4d99080c8e90ab40f0c27bda7a0f29aa1/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b32331d4f4df5d6eefa0554c565b626c7216f87a06a4f56fab27c3b68a830ec9", size = 63619976, upload-time = "2025-03-07T01:46:23.323Z" }, + { url = "https://files.pythonhosted.org/packages/1e/72/7c2ae24fb6b63a32e6ae5d241cc65263ea18d08802aaae087d9f013335a2/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:133df5a7509c3e292aaa2b477afd0194f06ce4ea24d714d616ff36439cee349a", size = 61962106, upload-time = "2025-08-04T10:21:41.128Z" }, + { url = "https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc", size = 59544258, upload-time = "2025-08-04T10:22:03.992Z" }, ] [[package]] -name = "nvidia-cusolver-cu12" -version = "11.7.3.90" +name = "nvidia-cusolver" +version = "12.0.4.66" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cublas", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cusparse", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" }, + { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, + { url = "https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112", size = 200941980, upload-time = "2025-09-04T08:33:22.767Z" }, ] [[package]] -name = "nvidia-cusparse-cu12" -version = "12.5.8.93" +name = "nvidia-cusparse" +version = "12.6.3.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" }, + { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, + { url = "https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b", size = 145942937, upload-time = "2025-09-04T08:33:58.029Z" }, ] [[package]] -name = "nvidia-cusparselt-cu12" -version = "0.7.1" +name = "nvidia-cusparselt-cu13" +version = "0.8.0" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" }, -] - -[[package]] -name = "nvidia-nccl-cu12" -version = "2.27.5" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/6e/89/f7a07dc961b60645dbbf42e80f2bc85ade7feb9a491b11a1e973aa00071f/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457", size = 322348229, upload-time = "2025-06-26T04:11:28.385Z" }, + { url = "https://files.pythonhosted.org/packages/46/10/8dcd1175260706a2fc92a16a52e306b71d4c1ea0b0cc4a9484183399818a/nvidia_cusparselt_cu13-0.8.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:400c6ed1cf6780fc6efedd64ec9f1345871767e6a1a0a552a1ea0578117ea77c", size = 220791277, upload-time = "2025-08-13T19:22:40.982Z" }, + { url = "https://files.pythonhosted.org/packages/fd/53/43b0d71f4e702fa9733f8b4571fdca50a8813f1e450b656c239beff12315/nvidia_cusparselt_cu13-0.8.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:25e30a8a7323935d4ad0340b95a0b69926eee755767e8e0b1cf8dd85b197d3fd", size = 169884119, upload-time = "2025-08-13T19:23:41.967Z" }, ] [[package]] name = "nvidia-nccl-cu13" -version = "2.29.2" +version = "2.28.9" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/cb/e8/b69bfcbf39d71b4166cf1ceb0e58dd73cc4c6ad005164b56e54acb4dbf2f/nvidia_nccl_cu13-2.29.2-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:9d4f7e24aff66309f0b71bd6a885afa214e1bac3a562c9a77be428f0a4aeb62a", size = 201038683, upload-time = "2026-01-07T00:21:18.002Z" }, - { url = "https://files.pythonhosted.org/packages/81/15/5e1d022945dd511d453ba5163fedce67d3bd0fb3dcadc021f00c0c8a491b/nvidia_nccl_cu13-2.29.2-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:86b997b315df0fb2874fd6062f2930d317bfa6434823351f287936d5ed616fc9", size = 201100704, upload-time = "2026-01-07T00:21:31.626Z" }, + { url = "https://files.pythonhosted.org/packages/39/55/1920646a2e43ffd4fc958536b276197ed740e9e0c54105b4bb3521591fc7/nvidia_nccl_cu13-2.28.9-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:01c873ba1626b54caa12272ed228dc5b2781545e0ae8ba3f432a8ef1c6d78643", size = 196561677, upload-time = "2025-11-18T05:49:03.45Z" }, + { url = "https://files.pythonhosted.org/packages/b0/b4/878fefaad5b2bcc6fcf8d474a25e3e3774bc5133e4b58adff4d0bca238bc/nvidia_nccl_cu13-2.28.9-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:e4553a30f34195f3fa1da02a6da3d6337d28f2003943aa0a3d247bbc25fefc42", size = 196493177, upload-time = "2025-11-18T05:49:17.677Z" }, ] [[package]] -name = "nvidia-nvjitlink-cu12" -version = "12.8.93" +name = "nvidia-nvjitlink" +version = "13.0.88" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f6/74/86a07f1d0f42998ca31312f998bd3b9a7eff7f52378f4f270c8679c77fb9/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88", size = 39254836, upload-time = "2025-03-07T01:49:55.661Z" }, + { url = "https://files.pythonhosted.org/packages/56/7a/123e033aaff487c77107195fa5a2b8686795ca537935a24efae476c41f05/nvidia_nvjitlink-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:13a74f429e23b921c1109976abefacc69835f2f433ebd323d3946e11d804e47b", size = 40713933, upload-time = "2025-09-04T08:35:43.553Z" }, + { url = "https://files.pythonhosted.org/packages/ab/2c/93c5250e64df4f894f1cbb397c6fd71f79813f9fd79d7cd61de3f97b3c2d/nvidia_nvjitlink-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e931536ccc7d467a98ba1d8b89ff7fa7f1fa3b13f2b0069118cd7f47bff07d0c", size = 38768748, upload-time = "2025-09-04T08:35:20.008Z" }, ] [[package]] -name = "nvidia-nvshmem-cu12" -version = "3.3.20" +name = "nvidia-nvshmem-cu13" +version = "3.4.5" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3b/6c/99acb2f9eb85c29fc6f3a7ac4dccfd992e22666dd08a642b303311326a97/nvidia_nvshmem_cu12-3.3.20-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d00f26d3f9b2e3c3065be895e3059d6479ea5c638a3f38c9fec49b1b9dd7c1e5", size = 124657145, upload-time = "2025-08-04T20:25:19.995Z" }, + { url = "https://files.pythonhosted.org/packages/dc/0f/05cc9c720236dcd2db9c1ab97fff629e96821be2e63103569da0c9b72f19/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dc2a197f38e5d0376ad52cd1a2a3617d3cdc150fd5966f4aee9bcebb1d68fe9", size = 60215947, upload-time = "2025-09-06T00:32:20.022Z" }, + { url = "https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:290f0a2ee94c9f3687a02502f3b9299a9f9fe826e6d0287ee18482e78d495b80", size = 60412546, upload-time = "2025-09-06T00:32:41.564Z" }, ] [[package]] -name = "nvidia-nvtx-cu12" -version = "12.8.90" +name = "nvidia-nvtx" +version = "13.0.85" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" }, + { url = "https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4", size = 148047, upload-time = "2025-09-04T08:29:01.761Z" }, + { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" }, ] [[package]] @@ -2918,46 +3012,37 @@ wheels = [ [[package]] name = "torch" -version = "2.9.1" +version = "2.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "cuda-bindings", marker = "sys_platform == 'linux'" }, + { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" }, { name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "fsspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "networkx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cufile-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cusparselt-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nvshmem-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" }, { name = "setuptools", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "sympy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "triton", marker = "sys_platform == 'linux'" }, { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/20/60/8fc5e828d050bddfab469b3fe78e5ab9a7e53dda9c3bdc6a43d17ce99e63/torch-2.9.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:c29455d2b910b98738131990394da3e50eea8291dfeb4b12de71ecf1fdeb21cb", size = 104135743, upload-time = "2025-11-12T15:21:34.936Z" }, - { url = "https://files.pythonhosted.org/packages/f2/b7/6d3f80e6918213babddb2a37b46dbb14c15b14c5f473e347869a51f40e1f/torch-2.9.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:524de44cd13931208ba2c4bde9ec7741fd4ae6bfd06409a604fc32f6520c2bc9", size = 899749493, upload-time = "2025-11-12T15:24:36.356Z" }, - { url = "https://files.pythonhosted.org/packages/28/0e/2a37247957e72c12151b33a01e4df651d9d155dd74d8cfcbfad15a79b44a/torch-2.9.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:5be4bf7496f1e3ffb1dd44b672adb1ac3f081f204c5ca81eba6442f5f634df8e", size = 74830751, upload-time = "2025-11-12T15:21:43.792Z" }, - { url = "https://files.pythonhosted.org/packages/4b/f7/7a18745edcd7b9ca2381aa03353647bca8aace91683c4975f19ac233809d/torch-2.9.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:30a3e170a84894f3652434b56d59a64a2c11366b0ed5776fab33c2439396bf9a", size = 104142929, upload-time = "2025-11-12T15:21:48.319Z" }, - { url = "https://files.pythonhosted.org/packages/f4/dd/f1c0d879f2863ef209e18823a988dc7a1bf40470750e3ebe927efdb9407f/torch-2.9.1-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:8301a7b431e51764629208d0edaa4f9e4c33e6df0f2f90b90e261d623df6a4e2", size = 899748978, upload-time = "2025-11-12T15:23:04.568Z" }, - { url = "https://files.pythonhosted.org/packages/40/60/71c698b466dd01e65d0e9514b5405faae200c52a76901baf6906856f17e4/torch-2.9.1-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:2c14b3da5df416cf9cb5efab83aa3056f5b8cd8620b8fde81b4987ecab730587", size = 74480347, upload-time = "2025-11-12T15:21:57.648Z" }, - { url = "https://files.pythonhosted.org/packages/48/50/c4b5112546d0d13cc9eaa1c732b823d676a9f49ae8b6f97772f795874a03/torch-2.9.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1edee27a7c9897f4e0b7c14cfc2f3008c571921134522d5b9b5ec4ebbc69041a", size = 74433245, upload-time = "2025-11-12T15:22:39.027Z" }, - { url = "https://files.pythonhosted.org/packages/81/c9/2628f408f0518b3bae49c95f5af3728b6ab498c8624ab1e03a43dd53d650/torch-2.9.1-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:19d144d6b3e29921f1fc70503e9f2fc572cde6a5115c0c0de2f7ca8b1483e8b6", size = 104134804, upload-time = "2025-11-12T15:22:35.222Z" }, - { url = "https://files.pythonhosted.org/packages/28/fc/5bc91d6d831ae41bf6e9e6da6468f25330522e92347c9156eb3f1cb95956/torch-2.9.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:c432d04376f6d9767a9852ea0def7b47a7bbc8e7af3b16ac9cf9ce02b12851c9", size = 899747132, upload-time = "2025-11-12T15:23:36.068Z" }, - { url = "https://files.pythonhosted.org/packages/bd/b2/2d15a52516b2ea3f414643b8de68fa4cb220d3877ac8b1028c83dc8ca1c4/torch-2.9.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:cb10896a1f7fedaddbccc2017ce6ca9ecaaf990f0973bdfcf405439750118d2c", size = 74823558, upload-time = "2025-11-12T15:22:43.392Z" }, - { url = "https://files.pythonhosted.org/packages/86/5c/5b2e5d84f5b9850cd1e71af07524d8cbb74cba19379800f1f9f7c997fc70/torch-2.9.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:0a2bd769944991c74acf0c4ef23603b9c777fdf7637f115605a4b2d8023110c7", size = 104145788, upload-time = "2025-11-12T15:23:52.109Z" }, - { url = "https://files.pythonhosted.org/packages/a9/8c/3da60787bcf70add986c4ad485993026ac0ca74f2fc21410bc4eb1bb7695/torch-2.9.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:07c8a9660bc9414c39cac530ac83b1fb1b679d7155824144a40a54f4a47bfa73", size = 899735500, upload-time = "2025-11-12T15:24:08.788Z" }, + { url = "https://files.pythonhosted.org/packages/87/89/5ea6722763acee56b045435fb84258db7375c48165ec8be7880ab2b281c5/torch-2.11.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1e6debd97ccd3205bbb37eb806a9d8219e1139d15419982c09e23ef7d4369d18", size = 80606801, upload-time = "2026-03-23T18:10:18.649Z" }, + { url = "https://files.pythonhosted.org/packages/32/d1/8ed2173589cbfe744ed54e5a73efc107c0085ba5777ee93a5f4c1ab90553/torch-2.11.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:63a68fa59de8f87acc7e85a5478bb2dddbb3392b7593ec3e78827c793c4b73fd", size = 419732382, upload-time = "2026-03-23T18:08:30.835Z" }, + { url = "https://files.pythonhosted.org/packages/3d/e1/b73f7c575a4b8f87a5928f50a1e35416b5e27295d8be9397d5293e7e8d4c/torch-2.11.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:cc89b9b173d9adfab59fd227f0ab5e5516d9a52b658ae41d64e59d2e55a418db", size = 530711509, upload-time = "2026-03-23T18:08:47.213Z" }, + { url = "https://files.pythonhosted.org/packages/db/38/8ac78069621b8c2b4979c2f96dc8409ef5e9c4189f6aac629189a78677ca/torch-2.11.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:8b394322f49af4362d4f80e424bcaca7efcd049619af03a4cf4501520bdf0fb4", size = 80959574, upload-time = "2026-03-23T18:10:14.214Z" }, + { url = "https://files.pythonhosted.org/packages/6d/6c/56bfb37073e7136e6dd86bfc6af7339946dd684e0ecf2155ac0eee687ae1/torch-2.11.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:2658f34ce7e2dabf4ec73b45e2ca68aedad7a5be87ea756ad656eaf32bf1e1ea", size = 419732324, upload-time = "2026-03-23T18:09:36.604Z" }, + { url = "https://files.pythonhosted.org/packages/07/f4/1b666b6d61d3394cca306ea543ed03a64aad0a201b6cd159f1d41010aeb1/torch-2.11.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:98bb213c3084cfe176302949bdc360074b18a9da7ab59ef2edc9d9f742504778", size = 530596026, upload-time = "2026-03-23T18:09:20.842Z" }, + { url = "https://files.pythonhosted.org/packages/26/0d/8603382f61abd0db35841148ddc1ffd607bf3100b11c6e1dab6d2fc44e72/torch-2.11.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:01018087326984a33b64e04c8cb5c2795f9120e0d775ada1f6638840227b04d7", size = 80573442, upload-time = "2026-03-23T18:09:10.117Z" }, + { url = "https://files.pythonhosted.org/packages/c7/86/7cd7c66cb9cec6be330fff36db5bd0eef386d80c031b581ec81be1d4b26c/torch-2.11.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:2bb3cc54bd0dea126b0060bb1ec9de0f9c7f7342d93d436646516b0330cd5be7", size = 419749385, upload-time = "2026-03-23T18:07:33.77Z" }, + { url = "https://files.pythonhosted.org/packages/47/e8/b98ca2d39b2e0e4730c0ee52537e488e7008025bc77ca89552ff91021f7c/torch-2.11.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:4dc8b3809469b6c30b411bb8c4cad3828efd26236153d9beb6a3ec500f211a60", size = 530716756, upload-time = "2026-03-23T18:07:50.02Z" }, + { url = "https://files.pythonhosted.org/packages/bf/46/4419098ed6d801750f26567b478fc185c3432e11e2cad712bc6b4c2ab0d0/torch-2.11.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:8245477871c3700d4370352ffec94b103cfcb737229445cf9946cddb7b2ca7cd", size = 80959460, upload-time = "2026-03-23T18:09:00.818Z" }, + { url = "https://files.pythonhosted.org/packages/fd/66/54a56a4a6ceaffb567231994a9745821d3af922a854ed33b0b3a278e0a99/torch-2.11.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:ab9a8482f475f9ba20e12db84b0e55e2f58784bdca43a854a6ccd3fd4b9f75e6", size = 419735835, upload-time = "2026-03-23T18:07:18.974Z" }, + { url = "https://files.pythonhosted.org/packages/b1/e7/0b6665f533aa9e337662dc190425abc0af1fe3234088f4454c52393ded61/torch-2.11.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:563ed3d25542d7e7bbc5b235ccfacfeb97fb470c7fee257eae599adb8005c8a2", size = 530613405, upload-time = "2026-03-23T18:08:07.014Z" }, ] [[package]] @@ -2971,10 +3056,9 @@ wheels = [ [[package]] name = "transformers" -version = "5.0.0" +version = "5.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "packaging", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, @@ -2983,22 +3067,26 @@ dependencies = [ { name = "safetensors", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "tokenizers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "tqdm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "typer-slim", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "typer", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/bc/79/845941711811789c85fb7e2599cea425a14a07eda40f50896b9d3fda7492/transformers-5.0.0.tar.gz", hash = "sha256:5f5634efed6cf76ad068cc5834c7adbc32db78bbd6211fb70df2325a9c37dec8", size = 8424830, upload-time = "2026-01-26T10:46:46.813Z" } +sdist = { url = "https://files.pythonhosted.org/packages/fc/1a/70e830d53ecc96ce69cfa8de38f163712d2b43ac52fbd743f39f56025c31/transformers-5.3.0.tar.gz", hash = "sha256:009555b364029da9e2946d41f1c5de9f15e6b1df46b189b7293f33a161b9c557", size = 8830831, upload-time = "2026-03-04T17:41:46.119Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/52/f3/ac976fa8e305c9e49772527e09fbdc27cc6831b8a2f6b6063406626be5dd/transformers-5.0.0-py3-none-any.whl", hash = "sha256:587086f249ce64c817213cf36afdb318d087f790723e9b3d4500b97832afd52d", size = 10142091, upload-time = "2026-01-26T10:46:43.88Z" }, + { url = "https://files.pythonhosted.org/packages/b8/88/ae8320064e32679a5429a2c9ebbc05c2bf32cefb6e076f9b07f6d685a9b4/transformers-5.3.0-py3-none-any.whl", hash = "sha256:50ac8c89c3c7033444fb3f9f53138096b997ebb70d4b5e50a2e810bf12d3d29a", size = 10661827, upload-time = "2026-03-04T17:41:42.722Z" }, ] [[package]] name = "triton" -version = "3.5.1" +version = "3.6.0" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/27/46/8c3bbb5b0a19313f50edcaa363b599e5a1a5ac9683ead82b9b80fe497c8d/triton-3.5.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f3f4346b6ebbd4fad18773f5ba839114f4826037c9f2f34e0148894cd5dd3dba", size = 170470410, upload-time = "2025-11-11T17:41:06.319Z" }, - { url = "https://files.pythonhosted.org/packages/37/92/e97fcc6b2c27cdb87ce5ee063d77f8f26f19f06916aa680464c8104ef0f6/triton-3.5.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0b4d2c70127fca6a23e247f9348b8adde979d2e7a20391bfbabaac6aebc7e6a8", size = 170579924, upload-time = "2025-11-11T17:41:12.455Z" }, - { url = "https://files.pythonhosted.org/packages/a4/e6/c595c35e5c50c4bc56a7bac96493dad321e9e29b953b526bbbe20f9911d0/triton-3.5.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d0637b1efb1db599a8e9dc960d53ab6e4637db7d4ab6630a0974705d77b14b60", size = 170480488, upload-time = "2025-11-11T17:41:18.222Z" }, - { url = "https://files.pythonhosted.org/packages/16/b5/b0d3d8b901b6a04ca38df5e24c27e53afb15b93624d7fd7d658c7cd9352a/triton-3.5.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bac7f7d959ad0f48c0e97d6643a1cc0fd5786fe61cb1f83b537c6b2d54776478", size = 170582192, upload-time = "2025-11-11T17:41:23.963Z" }, + { url = "https://files.pythonhosted.org/packages/3c/12/34d71b350e89a204c2c7777a9bba0dcf2f19a5bfdd70b57c4dbc5ffd7154/triton-3.6.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:448e02fe6dc898e9e5aa89cf0ee5c371e99df5aa5e8ad976a80b93334f3494fd", size = 176133521, upload-time = "2026-01-20T16:16:13.321Z" }, + { url = "https://files.pythonhosted.org/packages/f9/0b/37d991d8c130ce81a8728ae3c25b6e60935838e9be1b58791f5997b24a54/triton-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10c7f76c6e72d2ef08df639e3d0d30729112f47a56b0c81672edc05ee5116ac9", size = 188289450, upload-time = "2026-01-20T16:00:49.136Z" }, + { url = "https://files.pythonhosted.org/packages/ce/4e/41b0c8033b503fd3cfcd12392cdd256945026a91ff02452bef40ec34bee7/triton-3.6.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1722e172d34e32abc3eb7711d0025bb69d7959ebea84e3b7f7a341cd7ed694d6", size = 176276087, upload-time = "2026-01-20T16:16:18.989Z" }, + { url = "https://files.pythonhosted.org/packages/35/f8/9c66bfc55361ec6d0e4040a0337fb5924ceb23de4648b8a81ae9d33b2b38/triton-3.6.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d002e07d7180fd65e622134fbd980c9a3d4211fb85224b56a0a0efbd422ab72f", size = 188400296, upload-time = "2026-01-20T16:00:56.042Z" }, + { url = "https://files.pythonhosted.org/packages/49/55/5ecf0dcaa0f2fbbd4420f7ef227ee3cb172e91e5fede9d0ecaddc43363b4/triton-3.6.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ef5523241e7d1abca00f1d240949eebdd7c673b005edbbce0aca95b8191f1d43", size = 176138577, upload-time = "2026-01-20T16:16:25.426Z" }, + { url = "https://files.pythonhosted.org/packages/df/3d/9e7eee57b37c80cec63322c0231bb6da3cfe535a91d7a4d64896fcb89357/triton-3.6.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a17a5d5985f0ac494ed8a8e54568f092f7057ef60e1b0fa09d3fd1512064e803", size = 188273063, upload-time = "2026-01-20T16:01:07.278Z" }, + { url = "https://files.pythonhosted.org/packages/48/db/56ee649cab5eaff4757541325aca81f52d02d4a7cd3506776cad2451e060/triton-3.6.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0b3a97e8ed304dfa9bd23bb41ca04cdf6b2e617d5e782a8653d616037a5d537d", size = 176274804, upload-time = "2026-01-20T16:16:31.528Z" }, + { url = "https://files.pythonhosted.org/packages/f6/56/6113c23ff46c00aae423333eb58b3e60bdfe9179d542781955a5e1514cb3/triton-3.6.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:46bd1c1af4b6704e554cad2eeb3b0a6513a980d470ccfa63189737340c7746a7", size = 188397994, upload-time = "2026-01-20T16:01:14.236Z" }, ] [[package]] @@ -3041,16 +3129,18 @@ datetime = [ ] [[package]] -name = "typer-slim" -version = "0.21.1" +name = "typer" +version = "0.24.1" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "annotated-doc", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "click", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "rich", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "shellingham", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/17/d4/064570dec6358aa9049d4708e4a10407d74c99258f8b2136bb8702303f1a/typer_slim-0.21.1.tar.gz", hash = "sha256:73495dd08c2d0940d611c5a8c04e91c2a0a98600cbd4ee19192255a233b6dbfd", size = 110478, upload-time = "2026-01-06T11:21:11.176Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f5/24/cb09efec5cc954f7f9b930bf8279447d24618bb6758d4f6adf2574c41780/typer-0.24.1.tar.gz", hash = "sha256:e39b4732d65fbdcde189ae76cf7cd48aeae72919dea1fdfc16593be016256b45", size = 118613, upload-time = "2026-02-21T16:54:40.609Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c8/0a/4aca634faf693e33004796b6cee0ae2e1dba375a800c16ab8d3eff4bb800/typer_slim-0.21.1-py3-none-any.whl", hash = "sha256:6e6c31047f171ac93cc5a973c9e617dbc5ab2bddc4d0a3135dc161b4e2020e0d", size = 47444, upload-time = "2026-01-06T11:21:12.441Z" }, + { url = "https://files.pythonhosted.org/packages/4a/91/48db081e7a63bb37284f9fbcefda7c44c277b18b0e13fbc36ea2335b71e6/typer-0.24.1-py3-none-any.whl", hash = "sha256:112c1f0ce578bfb4cab9ffdabc68f031416ebcc216536611ba21f04e9aa84c9e", size = 56085, upload-time = "2026-02-21T16:54:41.616Z" }, ] [[package]] @@ -3092,6 +3182,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl", hash = "sha256:bf272323e553dfb2e87d9bfd225ca7b0f467b919d7bbd355436d3fd37cb0acd4", size = 131584, upload-time = "2026-01-07T16:24:42.685Z" }, ] +[[package]] +name = "uvicorn" +version = "0.42.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "h11", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e3/ad/4a96c425be6fb67e0621e62d86c402b4a17ab2be7f7c055d9bd2f638b9e2/uvicorn-0.42.0.tar.gz", hash = "sha256:9b1f190ce15a2dd22e7758651d9b6d12df09a13d51ba5bf4fc33c383a48e1775", size = 85393, upload-time = "2026-03-16T06:19:50.077Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0a/89/f8827ccff89c1586027a105e5630ff6139a64da2515e24dafe860bd9ae4d/uvicorn-0.42.0-py3-none-any.whl", hash = "sha256:96c30f5c7abe6f74ae8900a70e92b85ad6613b745d4879eb9b16ccad15645359", size = 68830, upload-time = "2026-03-16T06:19:48.325Z" }, +] + [[package]] name = "word2number" version = "1.1"