cortexgrid-infer 0.2.4__tar.gz → 0.2.7__tar.gz

{cortexgrid_infer-0.2.4 → cortexgrid_infer-0.2.7}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.5
 Name: cortexgrid-infer
-Version: 0.2.4
-Summary: Deploy and interact with models across providers (Anthropic, HuggingFace LLM + diffusion)
+Version: 0.2.7
+Summary: Deploy and interact with models across providers (Anthropic, HuggingFace LLM + diffusion + image-to-mesh)
 Project-URL: Homepage, https://github.com/robodatalab/cortexgrid-infer
 Project-URL: Repository, https://github.com/robodatalab/cortexgrid-infer
 License-Expression: Apache-2.0
@@ -14,6 +14,7 @@ Requires-Dist: diffusers>=0.38
 Requires-Dist: fastapi>=0.133
 Requires-Dist: httpx>=0.27.0
 Requires-Dist: huggingface-hub>=0.20.0
+Requires-Dist: numpy>=1.24
 Requires-Dist: pillow>=10.0
 Requires-Dist: protobuf>=4.0
 Requires-Dist: python-dotenv>=1.0.0
@@ -60,6 +61,7 @@ An **importer** is the whole of this library's cluster-side surface. It gives yo
 | `HuggingFaceCompletingImport(hf_id, token=None)` | HuggingFace causal LM, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
 | `HuggingFaceImageImport(hf_id, token=None, ignore_patterns=None)` | HuggingFace diffusers pipeline | `HuggingFaceImageModel` (`generate`) |
 | `AnthropicImport(model_id, api_key_secret=...)` | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
+| `HuggingFaceMeshImport(hf_id, token=None, ignore_patterns=None)` — a base: subclass it, naming a `MeshingDeployment` subclass as `serve_app` | an image-to-mesh model, run by the serve app the subclass supplies (`load` and `make_mesh`) behind the family's `POST /mesh` | `ServedMeshingModel` (`mesh`: one picture in, a `GeneratedMesh` out) |
 
 Not every model has weights, so `.source` belongs to the HuggingFace importers
 rather than to all of them: one that has weights goes through
@@ -184,16 +186,25 @@ Both serve apps are tuned rather than defaulted:
   itself, capturing prefill and decode separately, which it is better placed to
   do; compiling the module here as well would compile the same forward twice.
 
-  The fixed length is the point, and it comes from the model: the cache is
-  pinned to `config.max_position_embeddings`. A prompt longer than that is
-  truncated to it — keeping the end, since a chat template puts the turn to
-  answer last — with a warning naming both lengths. Serving it instead would
-  resize the cache and recompile the decode loop, which costs more than the
-  generation itself.
-
-Neither asks to be configured, for the same reason the serve apps do not ask which
-dtype to load in. Both happen on CUDA and nowhere else — inductor is weakest off
-it, and a host outrunning its accelerator is a GPU problem to begin with.
+  The fixed length is the point, and how long it is decides what decode costs:
+  every token's attention reads the whole cache, filled or not. That is why it is
+  *not* pinned to the model's context — 32k slots of a Qwen2.5-3B KV cache is
+  1.2 GB of keys and values re-read per generated token, several times the
+  traffic of the model's own weights, to make room for a prompt nobody sent. It
+  is pinned instead to the prompt and reply a caller plausibly sends: 4096 tokens,
+  or `max_total_tokens` on the model card, never past the model's own
+  `max_position_embeddings`. A reply takes at most half of it, so an
+  over-optimistic `max_new_tokens` cannot squeeze the prompt to nothing, and a
+  prompt that will not fit in the rest is truncated — keeping the end, since a
+  chat template puts the turn to answer last — with a warning naming both
+  lengths. Serving it whole would resize the cache and recompile the decode loop,
+  which costs more than the generation itself.
+
+The image app asks to be configured for nothing, for the same reason neither asks
+which dtype to load in; the completion app takes `max_total_tokens` alone, because
+how much cache a caller needs is the one thing the model cannot tell it. Both
+happen on CUDA and nowhere else — inductor is weakest off it, and a host
+outrunning its accelerator is a GPU problem to begin with.
 
 Two things worth knowing:
 

{cortexgrid_infer-0.2.4 → cortexgrid_infer-0.2.7}/README.md

@@ -34,6 +34,7 @@ An **importer** is the whole of this library's cluster-side surface. It gives yo
 | `HuggingFaceCompletingImport(hf_id, token=None)` | HuggingFace causal LM, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
 | `HuggingFaceImageImport(hf_id, token=None, ignore_patterns=None)` | HuggingFace diffusers pipeline | `HuggingFaceImageModel` (`generate`) |
 | `AnthropicImport(model_id, api_key_secret=...)` | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
+| `HuggingFaceMeshImport(hf_id, token=None, ignore_patterns=None)` — a base: subclass it, naming a `MeshingDeployment` subclass as `serve_app` | an image-to-mesh model, run by the serve app the subclass supplies (`load` and `make_mesh`) behind the family's `POST /mesh` | `ServedMeshingModel` (`mesh`: one picture in, a `GeneratedMesh` out) |
 
 Not every model has weights, so `.source` belongs to the HuggingFace importers
 rather than to all of them: one that has weights goes through
@@ -158,16 +159,25 @@ Both serve apps are tuned rather than defaulted:
   itself, capturing prefill and decode separately, which it is better placed to
   do; compiling the module here as well would compile the same forward twice.
 
-  The fixed length is the point, and it comes from the model: the cache is
-  pinned to `config.max_position_embeddings`. A prompt longer than that is
-  truncated to it — keeping the end, since a chat template puts the turn to
-  answer last — with a warning naming both lengths. Serving it instead would
-  resize the cache and recompile the decode loop, which costs more than the
-  generation itself.
-
-Neither asks to be configured, for the same reason the serve apps do not ask which
-dtype to load in. Both happen on CUDA and nowhere else — inductor is weakest off
-it, and a host outrunning its accelerator is a GPU problem to begin with.
+  The fixed length is the point, and how long it is decides what decode costs:
+  every token's attention reads the whole cache, filled or not. That is why it is
+  *not* pinned to the model's context — 32k slots of a Qwen2.5-3B KV cache is
+  1.2 GB of keys and values re-read per generated token, several times the
+  traffic of the model's own weights, to make room for a prompt nobody sent. It
+  is pinned instead to the prompt and reply a caller plausibly sends: 4096 tokens,
+  or `max_total_tokens` on the model card, never past the model's own
+  `max_position_embeddings`. A reply takes at most half of it, so an
+  over-optimistic `max_new_tokens` cannot squeeze the prompt to nothing, and a
+  prompt that will not fit in the rest is truncated — keeping the end, since a
+  chat template puts the turn to answer last — with a warning naming both
+  lengths. Serving it whole would resize the cache and recompile the decode loop,
+  which costs more than the generation itself.
+
+The image app asks to be configured for nothing, for the same reason neither asks
+which dtype to load in; the completion app takes `max_total_tokens` alone, because
+how much cache a caller needs is the one thing the model cannot tell it. Both
+happen on CUDA and nowhere else — inductor is weakest off it, and a host
+outrunning its accelerator is a GPU problem to begin with.
 
 Two things worth knowing:
 

{cortexgrid_infer-0.2.4 → cortexgrid_infer-0.2.7}/cortexgrid_infer/__init__.py

@@ -26,13 +26,16 @@ from cortexgrid_infer.core import (
     DeployedModel,
     CompletingModel,
     GeneratedImage,
+    GeneratedMesh,
     GeneratingModel,
+    MeshingModel,
     Message,
     Tool,
     ToolCall,
     ToolSpec,
     complete,
     generate,
+    mesh,
 )
 from cortexgrid_infer.completion import ServedCompletingModel
 from cortexgrid_infer.device import detect_device
@@ -43,19 +46,25 @@ from cortexgrid_infer.providers.huggingface_image import (
     HuggingFaceImageImport,
     HuggingFaceImageModel,
 )
+from cortexgrid_infer.meshing import ServedMeshingModel
+from cortexgrid_infer.meshing_serve import MeshingDeployment
+from cortexgrid_infer.providers.huggingface_mesh import HuggingFaceMeshImport
 
 __all__ = [
     "CompletionChunk",
     "DeployedModel",
     "CompletingModel",
     "GeneratedImage",
+    "GeneratedMesh",
     "GeneratingModel",
+    "MeshingModel",
     "Message",
     "Tool",
     "ToolCall",
     "ToolSpec",
     "complete",
     "generate",
+    "mesh",
     "ModelDeployFailed",
     "detect_device",
     "split_model_id",
@@ -63,7 +72,10 @@ __all__ = [
     "HuggingFaceImport",
     "HuggingFaceCompletingImport",
     "HuggingFaceImageImport",
+    "HuggingFaceMeshImport",
     "AnthropicImport",
     "ServedCompletingModel",
     "HuggingFaceImageModel",
+    "ServedMeshingModel",
+    "MeshingDeployment",
 ]

{cortexgrid_infer-0.2.4 → cortexgrid_infer-0.2.7}/cortexgrid_infer/core.py

@@ -8,6 +8,8 @@ from functools import partial
 from collections.abc import AsyncIterator
 from typing import Any, Callable, Sequence
 
+import numpy as np
+
 
 @dataclass
 class ToolCall:
@@ -87,6 +89,26 @@ class GeneratingModel(DeployedModel):
     ) -> GeneratedImage: ...
 
 
+@dataclass
+class GeneratedMesh:
+    """One mesh: `vertices` (N, 3) float32, x to the right of the picture it was
+    made from, y up, z toward whoever looked at it; `faces` (M, 3) int32 indices
+    into them, counter-clockwise seen from outside; `colours` (N, 3) uint8, each
+    vertex's colour. Plus the resolved parameters the deployment used."""
+
+    vertices: np.ndarray
+    faces: np.ndarray
+    colours: np.ndarray
+    params: dict[str, Any] = field(default_factory=dict)
+
+
+class MeshingModel(DeployedModel):
+    """A deployed model that makes a mesh of the object in one picture."""
+
+    @abc.abstractmethod
+    async def mesh(self, image: bytes, **kwargs: Any) -> GeneratedMesh: ...
+
+
 async def complete(
     deployed_model: CompletingModel,
     messages: list[Message],
@@ -115,3 +137,8 @@ async def generate(
 ) -> GeneratedImage:
     """Single image-generation request to the deployed model."""
     return await deployed_model.generate(prompt, image=image, **kwargs)
+
+
+async def mesh(deployed_model: MeshingModel, image: bytes, **kwargs: Any) -> GeneratedMesh:
+    """Single image-to-mesh request to the deployed model."""
+    return await deployed_model.mesh(image, **kwargs)

cortexgrid_infer-0.2.7/cortexgrid_infer/meshing.py

@@ -0,0 +1,72 @@
+"""The wire protocol between a cortexgrid-served image-to-mesh app and its client.
+
+An image-to-mesh serve app answers `POST /mesh`: the request carries a picture
+of an object as a base64 PNG under `image`, plus whatever options the model
+takes (a marching-cubes resolution, say) as further fields. The reply carries
+the mesh as base64 little-endian arrays - `vertices` float32 x, y, z, `faces`
+int32 indices, `colours` uint8 r, g, b per vertex - in the frame
+`GeneratedMesh` documents, and `params`, what the model resolved.
+
+Every such app speaks it through `cortexgrid_infer.meshing_serve.MeshingDeployment`,
+so one client, `ServedMeshingModel`, serves every model of the family whoever
+made it.
+"""
+
+from __future__ import annotations
+
+import base64
+from dataclasses import dataclass
+from typing import Any
+
+import httpx
+import numpy as np
+
+from cortexgrid_infer.core import GeneratedMesh, MeshingModel
+
+# How each array travels: its type, and how many values make one row.
+VERTICES = ("<f4", 3)
+FACES = ("<i4", 3)
+COLOURS = ("u1", 3)
+
+
+def encode(array: np.ndarray, kind: tuple[str, int]) -> str:
+    """`array` as the protocol sends it."""
+    dtype, _ = kind
+    return base64.b64encode(np.ascontiguousarray(array, dtype=dtype).tobytes()).decode("ascii")
+
+
+def decode(data: str, kind: tuple[str, int]) -> np.ndarray:
+    """An array the protocol sent, one row per vertex or triangle."""
+    dtype, columns = kind
+    return np.frombuffer(base64.b64decode(data), dtype=dtype).reshape(-1, columns)
+
+
+@dataclass
+class ServedMeshingModel(MeshingModel):
+    """The client of any image-to-mesh serve app deployed at `url`."""
+
+    url: str
+    model_id: str
+
+    @property
+    def name(self) -> str:
+        return self.model_id
+
+    async def mesh(self, image: bytes, **options: Any) -> GeneratedMesh:
+        body: dict[str, Any] = {"image": base64.b64encode(image).decode("ascii"), **options}
+        # Meshing takes seconds; no client-side timeout.
+        async with httpx.AsyncClient(timeout=None) as client:
+            response = await client.post(f"{self.url}/mesh", json=body)
+            response.raise_for_status()
+            data = response.json()
+
+        return GeneratedMesh(
+            vertices=decode(data["vertices"], VERTICES),
+            faces=decode(data["faces"], FACES),
+            colours=decode(data["colours"], COLOURS),
+            params={
+                **data.get("params", {}),
+                "duration_s": data.get("duration_s"),
+                "model_id": self.model_id,
+            },
+        )

cortexgrid_infer-0.2.7/cortexgrid_infer/meshing_serve.py

@@ -0,0 +1,72 @@
+"""The base of every image-to-mesh serve app: the `POST /mesh` route and loading the
+model's files from the cortexgrid registry, around the two steps that differ per
+model.
+
+There is no one loader for image-to-mesh models the way `AutoModelForCausalLM`
+or a diffusers `model_index.json` serves a whole family of repos: each model
+ships its own code. So a model's serve app subclasses `MeshingDeployment` and
+supplies `load` and `make_mesh`; the subclass inherits the route, whose protocol
+is `cortexgrid_infer.meshing`'s. cortexgrid bundles the subclass's own file, so
+the model's code travels with it.
+
+    class MyMeshDeployment(MeshingDeployment):
+        def load(self, path, device): ...
+        def make_mesh(self, image, **options) -> GeneratedMesh: ...
+"""
+
+from __future__ import annotations
+
+import base64
+import io
+import time
+from pathlib import Path
+from typing import Any
+
+import cortexgrid
+from cortexgrid import serve
+from fastapi import FastAPI
+from PIL import Image, ImageOps
+import torch
+
+from cortexgrid_infer import meshing
+from cortexgrid_infer.core import GeneratedMesh
+from cortexgrid_infer.device import detect_device
+
+
+_app = FastAPI()
+
+
+@serve.ingress(_app)
+class MeshingDeployment:
+    """One replica of an image-to-mesh model."""
+
+    def __init__(self, family: str, suffix: str, run_name: str) -> None:
+        self.device = detect_device()
+        self.load(Path(cortexgrid.load_model(family, suffix, run_name)), self.device)
+
+    def load(self, path: Path, device: torch.device) -> None:
+        """Build the model from its files at `path`, on `device`."""
+        raise NotImplementedError
+
+    def make_mesh(self, image: Image.Image, **options: Any) -> GeneratedMesh:
+        """The mesh of the object in `image`, in the frame `GeneratedMesh`
+        documents, with the options the request carried; `params` says what the
+        model resolved them to."""
+        raise NotImplementedError
+
+    @_app.post("/mesh")
+    async def mesh(self, body: dict[str, Any]) -> dict[str, Any]:
+        image = ImageOps.exif_transpose(Image.open(io.BytesIO(base64.b64decode(body["image"]))))
+        options = {key: value for key, value in body.items() if key != "image" and value is not None}
+
+        t0 = time.time()
+        made = self.make_mesh(image, **options)
+        duration = time.time() - t0
+
+        return {
+            "vertices": meshing.encode(made.vertices, meshing.VERTICES),
+            "faces": meshing.encode(made.faces, meshing.FACES),
+            "colours": meshing.encode(made.colours, meshing.COLOURS),
+            "params": made.params,
+            "duration_s": round(duration, 2),
+        }

{cortexgrid_infer-0.2.4 → cortexgrid_infer-0.2.7}/cortexgrid_infer/providers/huggingface_complete_serve.py

@@ -29,8 +29,14 @@ log = logging.getLogger(__name__)
 _app = FastAPI()
 
 
-# Used only by a model whose config states no context length of its own.
-FALLBACK_MAX_INPUT_TOKENS = 4096
+# Prompt + reply a replica is sized for. The static cache is allocated at this
+# many slots and every decoded token reads all of them, filled or not: at
+# Qwen2.5-3B's 32k context that is 1.2 GB of keys and values re-read per token,
+# which costs far more than the prompt it makes room for. So the default is what
+# a caller plausibly sends rather than what the model could hold; a model that
+# needs more takes `max_total_tokens` on its model card, up to its own context.
+MAX_TOTAL_TOKENS_PARAM = "max_total_tokens"
+DEFAULT_MAX_TOTAL_TOKENS = 4096
 
 _RESERVED_BODY_KEYS = {
     "messages",
@@ -55,8 +61,12 @@ class HuggingFaceCompletingDeployment:
             str(path), torch_dtype=torch.float16
         )
         self._model.to(self._device)  # type: ignore
-        self._max_input_tokens = getattr(
-            self._model.config, "max_position_embeddings", FALLBACK_MAX_INPUT_TOKENS
+        settings = cortexgrid.model_config(family, suffix, run_name)
+        context = getattr(
+            self._model.config, "max_position_embeddings", DEFAULT_MAX_TOTAL_TOKENS
+        )
+        self._max_total_tokens = min(
+            int(settings.get(MAX_TOTAL_TOKENS_PARAM, DEFAULT_MAX_TOTAL_TOKENS)), context
         )
         # A static cache of fixed length is what makes decode compilable: it
         # pins the shape every forward sees, so the loop is captured once
@@ -66,27 +76,39 @@ class HuggingFaceCompletingDeployment:
         if self._compiled:
             config = self._model.generation_config
             config.cache_implementation = "static"
-            config.max_cache_len = self._max_input_tokens
+            config.max_cache_len = self._max_total_tokens
             config.compile_config = CompileConfig(mode=compiling.Mode.GRAPHED)
 
-    def _truncate(self, inputs: Any) -> Any:
-        """Cut an over-long prompt to what the model handles, keeping its end."""
+    def _room_for_the_reply(self, requested: int) -> int:
+        """How many tokens a reply may take, leaving the prompt the rest.
+
+        Half the budget at most, so a caller asking for more output than the
+        replica is sized for cannot squeeze the prompt down to nothing."""
+        return min(requested, self._max_total_tokens // 2)
+
+    def _truncate(self, inputs: Any, limit: int) -> Any:
+        """Cut an over-long prompt to ``limit`` tokens, keeping its end.
+
+        A prompt and reply that together outgrew the cache would reallocate it
+        and recompile the decode loop, which costs more than the generation."""
         length = inputs["input_ids"].shape[-1]
-        if length <= self._max_input_tokens:
+        if length <= limit:
             return inputs
         log.warning(
-            "truncating prompt from %d to %d tokens, the model's context",
-            length, self._max_input_tokens,
+            "truncating prompt from %d to %d tokens, the room this replica has",
+            length, limit,
         )
         for key, value in inputs.items():
-            inputs[key] = value[..., -self._max_input_tokens :]
+            inputs[key] = value[..., -limit:]
         return inputs
 
     @_app.post("/complete")
     async def complete(self, body: dict[str, Any]) -> StreamingResponse:
         messages = body["messages"]
         tools = body.get("tools")
-        max_new_tokens = body.get("max_new_tokens", 16 * 1024)
+        max_new_tokens = self._room_for_the_reply(
+            body.get("max_new_tokens", self._max_total_tokens)
+        )
         temperature = body.get("temperature", 0.7)
         top_p = body.get("top_p", 0.9)
         top_k = body.get("top_k", 50)
@@ -99,9 +121,10 @@ class HuggingFaceCompletingDeployment:
             tokenize=False,
             add_generation_prompt=True,
         )
-        inputs = self._truncate(self._tokenizer(prompt, return_tensors="pt")).to(
-            self._model.device
-        )
+        inputs = self._truncate(
+            self._tokenizer(prompt, return_tensors="pt"),
+            self._max_total_tokens - max_new_tokens,
+        ).to(self._model.device)
 
         queue: asyncio.Queue[str | None] = asyncio.Queue()
         error: list[BaseException] = []

cortexgrid_infer-0.2.7/cortexgrid_infer/providers/huggingface_mesh.py

@@ -0,0 +1,37 @@
+"""HuggingFace image-to-mesh provider: the importer an image-to-mesh model on the
+Hub is taken into the cortexgrid registry with, and the client of what it serves.
+
+Each such model brings its own code (see `cortexgrid_infer.meshing_serve`), so
+this is a base: a model's own importer subclasses it and names its serve app.
+
+    class MyMeshImport(HuggingFaceMeshImport):
+        serve_app = MyMeshDeployment   # a MeshingDeployment subclass
+        vram_gb = 6.0
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import ClassVar
+
+import cortexgrid
+
+from cortexgrid_infer.importing import HuggingFaceImport
+from cortexgrid_infer.meshing import ServedMeshingModel
+
+
+class HuggingFaceMeshImport(HuggingFaceImport):
+    """What `cortexgrid.import_model` needs to take an image-to-mesh model from
+    the Hub, less the serve app, which the subclass names."""
+
+    # What one replica needs at least, whatever the weights come to: meshing
+    # queries the model over a whole 3D grid, which takes memory the weights do
+    # not show.
+    vram_gb: ClassVar[float] = 0.0
+
+    def requirements(self) -> cortexgrid.ModelRequirements:
+        estimate = super().requirements()
+        return dataclasses.replace(estimate, vram_gb=max(estimate.vram_gb, self.vram_gb))
+
+    def client(self, url: str) -> ServedMeshingModel:
+        return ServedMeshingModel(url=url, model_id=self.hf_id)

{cortexgrid_infer-0.2.4 → cortexgrid_infer-0.2.7}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "cortexgrid-infer"
-version = "0.2.4"
-description = "Deploy and interact with models across providers (Anthropic, HuggingFace LLM + diffusion)"
+version = "0.2.7"
+description = "Deploy and interact with models across providers (Anthropic, HuggingFace LLM + diffusion + image-to-mesh)"
 readme = "README.md"
 license = "Apache-2.0"
 license-files = ["LICENSE"]
@@ -34,6 +34,8 @@ dependencies = [
     "accelerate>=0.33",
     "sentencepiece>=0.2",
     "protobuf>=4.0",
+    # Image-to-mesh: meshes travel as numpy arrays (core.GeneratedMesh).
+    "numpy>=1.24",
 ]
 
 [project.optional-dependencies]