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

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

@@ -1,7 +1,7 @@
 Metadata-Version: 2.5
 Name: cortexgrid-infer
-Version: 0.2.6
-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

{cortexgrid_infer-0.2.6 → 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

{cortexgrid_infer-0.2.6 → 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.6 → 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.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.6 → cortexgrid_infer-0.2.7}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "cortexgrid-infer"
-version = "0.2.6"
-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]