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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (35) hide show
  1. {cortexgrid_infer-0.2.7 → cortexgrid_infer-0.3.2}/PKG-INFO +130 -56
  2. {cortexgrid_infer-0.2.7 → cortexgrid_infer-0.3.2}/README.md +127 -54
  3. {cortexgrid_infer-0.2.7 → cortexgrid_infer-0.3.2}/cortexgrid_infer/__init__.py +40 -22
  4. cortexgrid_infer-0.3.2/cortexgrid_infer/imaging.py +74 -0
  5. cortexgrid_infer-0.3.2/cortexgrid_infer/importers/__init__.py +10 -0
  6. cortexgrid_infer-0.3.2/cortexgrid_infer/importers/base.py +99 -0
  7. cortexgrid_infer-0.3.2/cortexgrid_infer/importers/huggingface.py +116 -0
  8. {cortexgrid_infer-0.2.7 → cortexgrid_infer-0.3.2}/cortexgrid_infer/meshing.py +2 -2
  9. cortexgrid_infer-0.3.2/cortexgrid_infer/models/__init__.py +21 -0
  10. cortexgrid_infer-0.2.7/cortexgrid_infer/providers/anthropic_serve.py → cortexgrid_infer-0.3.2/cortexgrid_infer/models/anthropic.py +30 -10
  11. cortexgrid_infer-0.3.2/cortexgrid_infer/models/base.py +121 -0
  12. cortexgrid_infer-0.3.2/cortexgrid_infer/models/gemini/__init__.py +11 -0
  13. cortexgrid_infer-0.3.2/cortexgrid_infer/models/gemini/base.py +50 -0
  14. cortexgrid_infer-0.3.2/cortexgrid_infer/models/gemini/text2image.py +116 -0
  15. cortexgrid_infer-0.3.2/cortexgrid_infer/models/gemini/text2text.py +146 -0
  16. cortexgrid_infer-0.2.7/cortexgrid_infer/meshing_serve.py → cortexgrid_infer-0.3.2/cortexgrid_infer/models/image2mesh.py +19 -7
  17. cortexgrid_infer-0.2.7/cortexgrid_infer/providers/huggingface_image_serve.py → cortexgrid_infer-0.3.2/cortexgrid_infer/models/text2image.py +27 -4
  18. cortexgrid_infer-0.2.7/cortexgrid_infer/providers/huggingface_complete_serve.py → cortexgrid_infer-0.3.2/cortexgrid_infer/models/text2text.py +15 -2
  19. cortexgrid_infer-0.3.2/cortexgrid_infer/registry.py +102 -0
  20. {cortexgrid_infer-0.2.7 → cortexgrid_infer-0.3.2}/pyproject.toml +6 -3
  21. cortexgrid_infer-0.2.7/cortexgrid_infer/importing.py +0 -158
  22. cortexgrid_infer-0.2.7/cortexgrid_infer/providers/__init__.py +0 -0
  23. cortexgrid_infer-0.2.7/cortexgrid_infer/providers/anthropic.py +0 -65
  24. cortexgrid_infer-0.2.7/cortexgrid_infer/providers/huggingface_complete.py +0 -23
  25. cortexgrid_infer-0.2.7/cortexgrid_infer/providers/huggingface_image.py +0 -110
  26. cortexgrid_infer-0.2.7/cortexgrid_infer/providers/huggingface_mesh.py +0 -37
  27. cortexgrid_infer-0.2.7/cortexgrid_infer/requirements.py +0 -117
  28. {cortexgrid_infer-0.2.7 → cortexgrid_infer-0.3.2}/.gitignore +0 -0
  29. {cortexgrid_infer-0.2.7 → cortexgrid_infer-0.3.2}/LICENSE +0 -0
  30. {cortexgrid_infer-0.2.7 → cortexgrid_infer-0.3.2}/cortexgrid_infer/compiling.py +0 -0
  31. {cortexgrid_infer-0.2.7 → cortexgrid_infer-0.3.2}/cortexgrid_infer/completion.py +0 -0
  32. {cortexgrid_infer-0.2.7 → cortexgrid_infer-0.3.2}/cortexgrid_infer/core.py +0 -0
  33. {cortexgrid_infer-0.2.7 → cortexgrid_infer-0.3.2}/cortexgrid_infer/device.py +0 -0
  34. {cortexgrid_infer-0.2.7 → cortexgrid_infer-0.3.2}/cortexgrid_infer/py.typed +0 -0
  35. {cortexgrid_infer-0.2.7 → cortexgrid_infer-0.3.2}/cortexgrid_infer/utils.py +0 -0
@@ -1,7 +1,7 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: cortexgrid-infer
3
- Version: 0.2.7
4
- Summary: Deploy and interact with models across providers (Anthropic, HuggingFace LLM + diffusion + image-to-mesh)
3
+ Version: 0.3.2
4
+ Summary: Serve apps per task (text-to-text, text-to-image, image-to-mesh) and importers per source (HuggingFace) for running models on cortexgrid
5
5
  Project-URL: Homepage, https://github.com/robodatalab/cortexgrid-infer
6
6
  Project-URL: Repository, https://github.com/robodatalab/cortexgrid-infer
7
7
  License-Expression: Apache-2.0
@@ -12,6 +12,7 @@ Requires-Dist: anthropic>=0.40.0
12
12
  Requires-Dist: cortexgrid>=0.3.4
13
13
  Requires-Dist: diffusers>=0.38
14
14
  Requires-Dist: fastapi>=0.133
15
+ Requires-Dist: google-genai>=1.0
15
16
  Requires-Dist: httpx>=0.27.0
16
17
  Requires-Dist: huggingface-hub>=0.20.0
17
18
  Requires-Dist: numpy>=1.24
@@ -33,19 +34,49 @@ needs but cannot work out for itself, plus the clients that talk to what it depl
33
34
  cortexgrid stores whatever weights it is handed under whatever key it is given, and
34
35
  serves them with whatever class it was told to bundle. It deliberately knows nothing
35
36
  about where a model came from or what shape it is. This library supplies exactly that
36
- missing knowledge for one family of models at a time, and drives no lifecycle of its
37
- own: importing, deploying and deleting stay your calls against the cortexgrid SDK.
37
+ missing knowledge, and drives no lifecycle of its own: importing, deploying and
38
+ deleting stay your calls against the cortexgrid SDK.
39
+
40
+ Upgrading from 0.2.x? See [MIGRATION.md](https://github.com/robodatalab/cortexgrid-infer/blob/main/MIGRATION.md).
38
41
 
39
42
  ## Division of labour
40
43
 
41
44
  | | Owns |
42
45
  |---|---|
43
46
  | **cortexgrid** | Registry identity, weights, serve bundle, hardware placement, deployment, jobs. Never knows what kind of model it holds. |
44
- | **cortexgrid-infer** | Per family: how to name it, fetch it, run it, size it, talk to it. Never drives a lifecycle. |
47
+ | **cortexgrid-infer** | Per source: how to name a model, fetch it and size it. Per task: how to run it and talk to it. Never drives a lifecycle. |
45
48
  | **you** | Policy: which model, which requirements, when to import, deploy and tear down. |
46
49
 
47
- An **importer** is the whole of this library's cluster-side surface. It gives you what
48
- `cortexgrid.import_model` and `cortexgrid.deploy_model` need:
50
+ A model is two things, and the library keeps them apart:
51
+
52
+ - **What runs it** — a serve app in `cortexgrid_infer.models`, named for the task it
53
+ runs, not for where the model came from. It loads the weights, answers the task's
54
+ routes, and names the client that speaks them. One serve app runs every model of
55
+ its task, whichever source staged it.
56
+ - **Where its weights come from** — an importer in `cortexgrid_infer.importers`, one
57
+ per source. It names the model in the registry, downloads the weights, and sizes
58
+ them from the source's metadata. It is handed the serve app, so one importer takes
59
+ any model its source holds.
60
+
61
+ | Serve app | Task | Runs | Client |
62
+ |---|---|---|---|
63
+ | `Text2Text` | text to text | any causal LM in the transformers layout, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
64
+ | `Text2Image` | text (and image) to image | any diffusers pipeline | `ServedGeneratingModel` (`generate`) |
65
+ | `Image2Mesh` — a base: subclass it with the model's `load` and `make_mesh` | image to mesh | the model the subclass loads, behind the task's `POST /mesh` | `ServedMeshingModel` (`mesh`: one picture in, a `GeneratedMesh` out) |
66
+ | `AnthropicText2Text` | text to text | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
67
+ | `GeminiText2Text` | text to text | the Gemini API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
68
+ | `GeminiText2Image` | text to image | the Gemini API, forwarded from the cluster | `ServedGeneratingModel` (`generate`) |
69
+
70
+ | Importer | Source |
71
+ |---|---|
72
+ | `HuggingFaceImporter(hf_id, serve_app, token=None, ignore_patterns=None)` | the HuggingFace Hub |
73
+
74
+ A model hosted elsewhere has no weights, and so no importer: its serve app forwards
75
+ to it, and a `Hosted(model_id, serve_app, **settings)` entry registers it.
76
+
77
+ An importer and a `Hosted` entry are both a `ModelEntry`, which gives you what
78
+ `cortexgrid.import_model`, `cortexgrid.register_model` and `cortexgrid.deploy_model`
79
+ need:
49
80
 
50
81
  | | |
51
82
  |---|---|
@@ -54,20 +85,10 @@ An **importer** is the whole of this library's cluster-side surface. It gives yo
54
85
  | `.requirements()` | what one replica needs to be placed and to run |
55
86
  | `.config()` | settings the serve app reads at construction, if it needs any |
56
87
  | `.client(url)` | a client that speaks the deployed app's routes |
57
- | `.source` | *(models with weights)* a callable that downloads them and returns the directory |
58
-
59
- | Importer | Serves | Client |
60
- |---|---|---|
61
- | `HuggingFaceCompletingImport(hf_id, token=None)` | HuggingFace causal LM, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
62
- | `HuggingFaceImageImport(hf_id, token=None, ignore_patterns=None)` | HuggingFace diffusers pipeline | `HuggingFaceImageModel` (`generate`) |
63
- | `AnthropicImport(model_id, api_key_secret=...)` | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
64
- | `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) |
88
+ | `.source` | *(importers only)* a callable that downloads the weights and returns the directory |
65
89
 
66
- Not every model has weights, so `.source` belongs to the HuggingFace importers
67
- rather than to all of them: one that has weights goes through
68
- `cortexgrid.import_model`, one that has none through `cortexgrid.register_model`.
69
- Everything else is common, which is why one client serves all three — every
70
- completion app speaks the same `/complete` protocol.
90
+ One with weights goes through `cortexgrid.import_model`, one without through
91
+ `cortexgrid.register_model`.
71
92
 
72
93
  ## Quick start
73
94
 
@@ -77,7 +98,7 @@ import cortexgrid
77
98
  import cortexgrid_infer as mg
78
99
 
79
100
  cortexgrid.Experiment.init("img-gen")
80
- imp = mg.HuggingFaceImageImport("black-forest-labs/FLUX.2-klein-base-4B")
101
+ imp = mg.HuggingFaceImporter("black-forest-labs/FLUX.2-klein-base-4B", mg.Text2Image)
81
102
 
82
103
  # 1. Import — weights go HuggingFace ─▶ cluster node ─▶ registry, never through
83
104
  # your machine, so this runs as a job rather than in-process.
@@ -132,12 +153,14 @@ serve code if it changed, and tags the run with the model.
132
153
  `cortexgrid.ModelRequirements` decides which node Ray will place a replica on, and
133
154
  nothing in the registry can derive it — cortexgrid stores whatever weights it is
134
155
  handed, and an importer imports whatever repo id it is handed. `requirements()`
135
- estimates it from the repo's own metadata: safetensors headers where the repo has
136
- them, weight-file sizes otherwise. Both are read over HTTP; no weights are
137
- downloaded.
156
+ splits the estimate along the same line as the rest: the importer reads what the
157
+ source records about the weights — for HuggingFace, safetensors headers where the
158
+ repo has them, weight-file sizes otherwise, both over HTTP with no weights
159
+ downloaded — and the serve app turns that into what a replica needs, since it knows
160
+ the dtype it loads in and what its work takes on top.
138
161
 
139
162
  ```python
140
- >>> mg.HuggingFaceCompletingImport("Qwen/Qwen2.5-0.5B-Instruct").requirements()
163
+ >>> mg.HuggingFaceImporter("Qwen/Qwen2.5-0.5B-Instruct", mg.Text2Text).requirements()
141
164
  ModelRequirements(num_gpus=1, ram_gb=3.0, vram_gb=3.2)
142
165
  ```
143
166
 
@@ -146,14 +169,15 @@ you know better, or edit the figures on the model card in the dashboard afterwar
146
169
  `import_model` only fills in requirements a version does not already have, so a
147
170
  hand-set value is never overwritten.
148
171
 
149
- The estimate always asks for exactly one GPU, because both HuggingFace serve apps move
150
- the whole model onto a single device. A model too large for any one card needs a serve
151
- app that shards it, not a larger `num_gpus`: the requirement would simply never be
152
- placed. (`AnthropicImport` asks for none — it holds no weights.)
172
+ The estimate always asks for exactly one GPU, because every serve app that runs weights
173
+ moves the whole model onto a single device. A model too large for any one card needs a
174
+ serve app that shards it, not a larger `num_gpus`: the requirement would simply never
175
+ be placed. An `Image2Mesh` subclass can set `min_vram_gb` for the memory meshing takes
176
+ beyond the weights. (A `Hosted` entry asks for none — it holds no weights.)
153
177
 
154
178
  ## Compilation
155
179
 
156
- Both HuggingFace serve apps run `torch.compile` over the model they load, because
180
+ `Text2Text` and `Text2Image` run `torch.compile` over the model they load, because
157
181
  eager PyTorch bills per *operation* and a transformer forward is mostly glue —
158
182
  views, broadcasts, permutes, elementwise adds — with only a small fraction of the
159
183
  dispatches being the matrix multiplies that do the arithmetic. While the tensors
@@ -174,12 +198,12 @@ fits the work it does:
174
198
 
175
199
  Both serve apps are tuned rather than defaulted:
176
200
 
177
- - **`HuggingFaceImageDeployment`** compiles the denoiser (`transformer` or `unet`)
201
+ - **`Text2Image`** compiles the denoiser (`transformer` or `unet`)
178
202
  and every `text_encoder*` as `Mode.GRAPHED`. A denoising schedule is the ideal case —
179
203
  every step is the same shapes — so the capture is made once and replayed for the
180
204
  whole loop. The VAE is left alone: smallest win of the three, and the one whose
181
205
  shapes move most, especially with tiling on.
182
- - **`HuggingFaceCompletingDeployment`** does *not* compile the module. Decode has
206
+ - **`Text2Text`** does *not* compile the module. Decode has
183
207
  no stable shape of its own — the sequence grows a token per forward — so it
184
208
  pins one: `cache_implementation="static"` with a fixed `max_cache_len`, plus a
185
209
  `CompileConfig(mode=Mode.GRAPHED)`. transformers then compiles `generate`
@@ -223,7 +247,7 @@ it is the floor; only quantizing the weights or running fewer forwards moves it.
223
247
  `complete` streams `CompletionChunk`s (`.content`, `.tool_calls`, `.finish_reason`):
224
248
 
225
249
  ```python
226
- imp = mg.HuggingFaceCompletingImport("Qwen/Qwen2.5-7B-Instruct")
250
+ imp = mg.HuggingFaceImporter("Qwen/Qwen2.5-7B-Instruct", mg.Text2Text)
227
251
  model = imp.client(deployment.url)
228
252
  messages = [{"role": "user", "content": "Explain RAG in one sentence."}]
229
253
 
@@ -238,17 +262,17 @@ async for chunk in mg.complete(model, messages, max_new_tokens=512, temperature=
238
262
  ### Anthropic
239
263
 
240
264
  An Anthropic model is a registry entry like any other — the difference is that
241
- there are no weights to stage, so it is imported with no source, and the serve
242
- app forwards to the API instead of loading anything:
265
+ there are no weights to stage, so there is no importer: the serve app forwards to
266
+ the API instead of loading anything, and a `Hosted` entry registers it:
243
267
 
244
268
  ```python
245
269
  cortexgrid.set_secret("ANTHROPIC_API_KEY", "sk-ant-...") # once
246
270
 
247
- imp = mg.AnthropicImport("claude-sonnet-5")
271
+ entry = mg.Hosted("claude-sonnet-5", mg.AnthropicText2Text) # api_key_secret="..." to use another secret
248
272
  cortexgrid.register_model(
249
- imp.serve_app,
250
- family=imp.family, suffix=imp.suffix,
251
- requirements=imp.requirements(), config=imp.config(),
273
+ entry.serve_app,
274
+ family=entry.family, suffix=entry.suffix,
275
+ requirements=entry.requirements(), config=entry.config(),
252
276
  )
253
277
  ```
254
278
 
@@ -265,45 +289,95 @@ a re-deploy rather than a re-registration.
265
289
 
266
290
  A replica asks for no hardware at all, so it is placed on any node, CPU-only
267
291
  included. From there it deploys and streams exactly like a cluster-served model:
268
- `imp.client(deployment.url)` returns the same `ServedCompletingModel`. Anthropic
292
+ `entry.client(deployment.url)` returns the same `ServedCompletingModel`. Anthropic
269
293
  reports tool calls as structured blocks rather than as generated text, so the
270
294
  serve app re-encodes them into the text form the client parses.
271
295
 
296
+ ### Gemini
297
+
298
+ A Gemini model registers exactly like an Anthropic one, with the serve app for
299
+ its task — `GeminiText2Text` or `GeminiText2Image`. The model id passed to
300
+ `Hosted` is the Gemini model name the deployment calls:
301
+
302
+ ```python
303
+ cortexgrid.set_secret("GEMINI_API_KEY", "AIza...") # once
304
+
305
+ entry = mg.Hosted("gemini-2.5-pro", mg.GeminiText2Text) # or "gemini-2.5-flash", ...
306
+ cortexgrid.register_model(
307
+ entry.serve_app,
308
+ family=entry.family, suffix=entry.suffix,
309
+ requirements=entry.requirements(), config=entry.config(),
310
+ )
311
+ ```
312
+
313
+ Everything said above for Anthropic holds: `config` carries the model name and
314
+ the name of the secret (`api_key_secret="..."` to use another), both editable on
315
+ the model card; a replica needs no hardware; and `entry.client(deployment.url)`
316
+ is a `ServedCompletingModel`, with Gemini's function calls re-encoded into the
317
+ text form the client parses.
318
+
319
+ `GeminiText2Image` answers the same `/generate` route as `Text2Image`, so its
320
+ client is a `ServedGeneratingModel`:
321
+
322
+ ```python
323
+ entry = mg.Hosted("gemini-2.5-flash-image", mg.GeminiText2Image)
324
+ ...
325
+ picture = await mg.generate(entry.client(deployment.url), "a fox in the snow", size=2048)
326
+ ```
327
+
328
+ It takes a prompt only, and makes a square picture. `size` picks Gemini's size
329
+ tier (up to 1K, 2K or 4K) and `seed` is passed on; `steps` and `guidance` are
330
+ diffusion settings with no Gemini counterpart, and come back as `None`.
331
+
272
332
  ## API
273
333
 
274
334
  | | |
275
335
  |---|---|
276
- | `HuggingFaceCompletingImport(hf_id, token=None)` | Importer for a causal LM. |
277
- | `HuggingFaceImageImport(hf_id, token=None, ignore_patterns=None)` | Importer for a diffusers pipeline. |
278
- | `AnthropicImport(model_id, api_key_secret="ANTHROPIC_API_KEY")` | Importer for an Anthropic model; no weights. |
279
- | `ModelImport` | Base of all three: `family`, `suffix`, `serve_app`, `requirements()`, `config()`, `client(url)`. |
280
- | `HuggingFaceImport` | Adds the download half — `source` and the scratch directory. Subclass it for another HuggingFace family. |
336
+ | `HuggingFaceImporter(hf_id, serve_app, token=None, ignore_patterns=None)` | Importer for a model on the HuggingFace Hub, run by `serve_app`. |
337
+ | `Importer` | Base of the importers: identity, scratch directory, `source`, and `requirements()` / `client(url)` taken from the serve app. A new source implements `download(local_dir)` and `weights()`. |
338
+ | `Hosted(model_id, serve_app, **settings)` | Entry for a model hosted elsewhere; no weights. `settings` are the serve app's `config` arguments. |
339
+ | `ModelEntry` | Base of both: `model_id`, `family`, `suffix`, `serve_app`, `requirements()`, `config()`, `client(url)`. |
340
+ | `Text2Text`, `Text2Image`, `Image2Mesh` | Serve apps that run weights (`LocalModel`s). |
341
+ | `AnthropicText2Text` | Serve app that forwards to the Anthropic API (a `HostedModel`). |
342
+ | `GeminiText2Text`, `GeminiText2Image` | Serve apps that forward to the Gemini API (`HostedModel`s sharing the `GeminiModel` base). |
343
+ | `LocalModel` | Base of the serve apps that run weights: `ignore_patterns()`, `bytes_per_param`, `min_vram_gb`, `requirements(weights)`, `client(url, name)`. |
344
+ | `HostedModel` | Base of the serve apps that forward: `config(model_id, **settings)`, `requirements()`, `client(url, name)`. |
345
+ | `Weights(params=None, file_bytes=0)` | What an importer reads about the weights, handed to the serve app to size a replica. |
281
346
  | `complete(model, messages, tools=None, max_new_tokens=2048, temperature=0.7, **kw)` | Async stream of `CompletionChunk` for a `CompletingModel`. |
282
347
  | `generate(model, prompt, *, image=None, **kw) -> GeneratedImage` | One image from a `GeneratingModel`. |
283
348
  | `ServedCompletingModel(url, model_id)` | Client for any deployed completion app. |
284
- | `HuggingFaceImageModel(url, model_id)` | Client for a deployed diffusers app. |
349
+ | `ServedGeneratingModel(url, model_id)` | Client for any deployed text-to-image app. |
350
+ | `ServedMeshingModel(url, model_id)` | Client for any deployed image-to-mesh app. |
285
351
  | `split_model_id(model_id) -> (family, suffix)` | The registry identity an importer derives. |
286
352
  | `detect_device()` | The torch device a serve app should load onto. |
287
353
  | `compiling.Mode.GRAPHED` / `compiling.Mode.FUSED` | The two modes a serve app chooses between. See [Compilation](#compilation). |
288
354
  | `compiling.supported(device)` / `compile_module(module, device, mode)` / `compile_pipeline(pipe, device, mode)` | Whether to compile here, and what a serve app calls to do it. |
289
355
  | `ModelDeployFailed` | Re-exported from cortexgrid; subclasses `RuntimeError`. |
290
356
 
291
- Types: `DeployedModel`, `CompletingModel`, `GeneratingModel`, `CompletionChunk`,
292
- `GeneratedImage`, `ToolCall`, `Message`, `Tool`, `ToolSpec`.
357
+ Types: `DeployedModel`, `CompletingModel`, `GeneratingModel`, `MeshingModel`,
358
+ `CompletionChunk`, `GeneratedImage`, `GeneratedMesh`, `ToolCall`, `Message`, `Tool`,
359
+ `ToolSpec`.
360
+
361
+ Extending it:
293
362
 
294
- Adding a provider means one importer and one serve app. A completion app only has
295
- to stream text from `POST /complete`, inlining tool calls as
296
- `<tool_call>{"name": ..., "arguments": {...}}</tool_call>` (there is an
297
- `encode_tool_call` for upstreams that report them structurally) — then
298
- `ServedCompletingModel` is its client, unchanged.
363
+ - **A new source** is one importer: subclass `Importer` with `download(local_dir)`
364
+ and `weights()`, and every serve app whose weight format it delivers runs its
365
+ models unchanged.
366
+ - **A new hosted API** is one serve app: subclass `HostedModel`, and register it
367
+ with `Hosted`. A text-to-text one only has to stream text from `POST /complete`,
368
+ inlining tool calls as `<tool_call>{"name": ..., "arguments": {...}}</tool_call>`
369
+ (there is an `encode_tool_call` for upstreams that report them structurally) —
370
+ then `ServedCompletingModel` is its client, unchanged.
371
+ - **A model no task's serve app can load** gets its own serve app — an `Image2Mesh`
372
+ subclass, say — and, if its source is unusual too, its own importer.
299
373
 
300
374
  ## Notes
301
375
 
302
376
  - **HuggingFace auth.** For gated/private repos, pass `token=` to the importer (or
303
377
  read `HF_TOKEN` yourself and pass it); it travels with the importer to the node
304
378
  that downloads.
305
- - **Image weight pruning.** `HuggingFaceImageImport` skips example images, docs and
306
- `.gitattributes` by default — nothing `from_pretrained` reads. It does **not**
379
+ - **Image weight pruning.** `Text2Image` tells the importer to skip example images,
380
+ docs and `.gitattributes` — nothing `from_pretrained` reads. It does **not**
307
381
  skip a repo's consolidated single-file checkpoint, which is weights it cannot
308
382
  tell apart from the component ones: name it per-model in
309
383
  `HF_IMAGE_SNAPSHOT_IGNORE` (comma-separated globs, e.g.
@@ -313,7 +387,7 @@ to stream text from `POST /complete`, inlining tool calls as
313
387
  describe what actually gets staged:
314
388
 
315
389
  ```python
316
- >>> mg.HuggingFaceImageImport("black-forest-labs/FLUX.2-klein-base-4B").requirements()
390
+ >>> mg.HuggingFaceImporter("black-forest-labs/FLUX.2-klein-base-4B", mg.Text2Image).requirements()
317
391
  ModelRequirements(num_gpus=1, ram_gb=26.3, vram_gb=29.6) # single-file checkpoint included
318
392
  ```
319
393
  - **Only model files are stored.** HuggingFace's download bookkeeping
@@ -6,19 +6,49 @@ needs but cannot work out for itself, plus the clients that talk to what it depl
6
6
  cortexgrid stores whatever weights it is handed under whatever key it is given, and
7
7
  serves them with whatever class it was told to bundle. It deliberately knows nothing
8
8
  about where a model came from or what shape it is. This library supplies exactly that
9
- missing knowledge for one family of models at a time, and drives no lifecycle of its
10
- own: importing, deploying and deleting stay your calls against the cortexgrid SDK.
9
+ missing knowledge, and drives no lifecycle of its own: importing, deploying and
10
+ deleting stay your calls against the cortexgrid SDK.
11
+
12
+ Upgrading from 0.2.x? See [MIGRATION.md](https://github.com/robodatalab/cortexgrid-infer/blob/main/MIGRATION.md).
11
13
 
12
14
  ## Division of labour
13
15
 
14
16
  | | Owns |
15
17
  |---|---|
16
18
  | **cortexgrid** | Registry identity, weights, serve bundle, hardware placement, deployment, jobs. Never knows what kind of model it holds. |
17
- | **cortexgrid-infer** | Per family: how to name it, fetch it, run it, size it, talk to it. Never drives a lifecycle. |
19
+ | **cortexgrid-infer** | Per source: how to name a model, fetch it and size it. Per task: how to run it and talk to it. Never drives a lifecycle. |
18
20
  | **you** | Policy: which model, which requirements, when to import, deploy and tear down. |
19
21
 
20
- An **importer** is the whole of this library's cluster-side surface. It gives you what
21
- `cortexgrid.import_model` and `cortexgrid.deploy_model` need:
22
+ A model is two things, and the library keeps them apart:
23
+
24
+ - **What runs it** — a serve app in `cortexgrid_infer.models`, named for the task it
25
+ runs, not for where the model came from. It loads the weights, answers the task's
26
+ routes, and names the client that speaks them. One serve app runs every model of
27
+ its task, whichever source staged it.
28
+ - **Where its weights come from** — an importer in `cortexgrid_infer.importers`, one
29
+ per source. It names the model in the registry, downloads the weights, and sizes
30
+ them from the source's metadata. It is handed the serve app, so one importer takes
31
+ any model its source holds.
32
+
33
+ | Serve app | Task | Runs | Client |
34
+ |---|---|---|---|
35
+ | `Text2Text` | text to text | any causal LM in the transformers layout, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
36
+ | `Text2Image` | text (and image) to image | any diffusers pipeline | `ServedGeneratingModel` (`generate`) |
37
+ | `Image2Mesh` — a base: subclass it with the model's `load` and `make_mesh` | image to mesh | the model the subclass loads, behind the task's `POST /mesh` | `ServedMeshingModel` (`mesh`: one picture in, a `GeneratedMesh` out) |
38
+ | `AnthropicText2Text` | text to text | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
39
+ | `GeminiText2Text` | text to text | the Gemini API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
40
+ | `GeminiText2Image` | text to image | the Gemini API, forwarded from the cluster | `ServedGeneratingModel` (`generate`) |
41
+
42
+ | Importer | Source |
43
+ |---|---|
44
+ | `HuggingFaceImporter(hf_id, serve_app, token=None, ignore_patterns=None)` | the HuggingFace Hub |
45
+
46
+ A model hosted elsewhere has no weights, and so no importer: its serve app forwards
47
+ to it, and a `Hosted(model_id, serve_app, **settings)` entry registers it.
48
+
49
+ An importer and a `Hosted` entry are both a `ModelEntry`, which gives you what
50
+ `cortexgrid.import_model`, `cortexgrid.register_model` and `cortexgrid.deploy_model`
51
+ need:
22
52
 
23
53
  | | |
24
54
  |---|---|
@@ -27,20 +57,10 @@ An **importer** is the whole of this library's cluster-side surface. It gives yo
27
57
  | `.requirements()` | what one replica needs to be placed and to run |
28
58
  | `.config()` | settings the serve app reads at construction, if it needs any |
29
59
  | `.client(url)` | a client that speaks the deployed app's routes |
30
- | `.source` | *(models with weights)* a callable that downloads them and returns the directory |
31
-
32
- | Importer | Serves | Client |
33
- |---|---|---|
34
- | `HuggingFaceCompletingImport(hf_id, token=None)` | HuggingFace causal LM, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
35
- | `HuggingFaceImageImport(hf_id, token=None, ignore_patterns=None)` | HuggingFace diffusers pipeline | `HuggingFaceImageModel` (`generate`) |
36
- | `AnthropicImport(model_id, api_key_secret=...)` | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
37
- | `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) |
60
+ | `.source` | *(importers only)* a callable that downloads the weights and returns the directory |
38
61
 
39
- Not every model has weights, so `.source` belongs to the HuggingFace importers
40
- rather than to all of them: one that has weights goes through
41
- `cortexgrid.import_model`, one that has none through `cortexgrid.register_model`.
42
- Everything else is common, which is why one client serves all three — every
43
- completion app speaks the same `/complete` protocol.
62
+ One with weights goes through `cortexgrid.import_model`, one without through
63
+ `cortexgrid.register_model`.
44
64
 
45
65
  ## Quick start
46
66
 
@@ -50,7 +70,7 @@ import cortexgrid
50
70
  import cortexgrid_infer as mg
51
71
 
52
72
  cortexgrid.Experiment.init("img-gen")
53
- imp = mg.HuggingFaceImageImport("black-forest-labs/FLUX.2-klein-base-4B")
73
+ imp = mg.HuggingFaceImporter("black-forest-labs/FLUX.2-klein-base-4B", mg.Text2Image)
54
74
 
55
75
  # 1. Import — weights go HuggingFace ─▶ cluster node ─▶ registry, never through
56
76
  # your machine, so this runs as a job rather than in-process.
@@ -105,12 +125,14 @@ serve code if it changed, and tags the run with the model.
105
125
  `cortexgrid.ModelRequirements` decides which node Ray will place a replica on, and
106
126
  nothing in the registry can derive it — cortexgrid stores whatever weights it is
107
127
  handed, and an importer imports whatever repo id it is handed. `requirements()`
108
- estimates it from the repo's own metadata: safetensors headers where the repo has
109
- them, weight-file sizes otherwise. Both are read over HTTP; no weights are
110
- downloaded.
128
+ splits the estimate along the same line as the rest: the importer reads what the
129
+ source records about the weights — for HuggingFace, safetensors headers where the
130
+ repo has them, weight-file sizes otherwise, both over HTTP with no weights
131
+ downloaded — and the serve app turns that into what a replica needs, since it knows
132
+ the dtype it loads in and what its work takes on top.
111
133
 
112
134
  ```python
113
- >>> mg.HuggingFaceCompletingImport("Qwen/Qwen2.5-0.5B-Instruct").requirements()
135
+ >>> mg.HuggingFaceImporter("Qwen/Qwen2.5-0.5B-Instruct", mg.Text2Text).requirements()
114
136
  ModelRequirements(num_gpus=1, ram_gb=3.0, vram_gb=3.2)
115
137
  ```
116
138
 
@@ -119,14 +141,15 @@ you know better, or edit the figures on the model card in the dashboard afterwar
119
141
  `import_model` only fills in requirements a version does not already have, so a
120
142
  hand-set value is never overwritten.
121
143
 
122
- The estimate always asks for exactly one GPU, because both HuggingFace serve apps move
123
- the whole model onto a single device. A model too large for any one card needs a serve
124
- app that shards it, not a larger `num_gpus`: the requirement would simply never be
125
- placed. (`AnthropicImport` asks for none — it holds no weights.)
144
+ The estimate always asks for exactly one GPU, because every serve app that runs weights
145
+ moves the whole model onto a single device. A model too large for any one card needs a
146
+ serve app that shards it, not a larger `num_gpus`: the requirement would simply never
147
+ be placed. An `Image2Mesh` subclass can set `min_vram_gb` for the memory meshing takes
148
+ beyond the weights. (A `Hosted` entry asks for none — it holds no weights.)
126
149
 
127
150
  ## Compilation
128
151
 
129
- Both HuggingFace serve apps run `torch.compile` over the model they load, because
152
+ `Text2Text` and `Text2Image` run `torch.compile` over the model they load, because
130
153
  eager PyTorch bills per *operation* and a transformer forward is mostly glue —
131
154
  views, broadcasts, permutes, elementwise adds — with only a small fraction of the
132
155
  dispatches being the matrix multiplies that do the arithmetic. While the tensors
@@ -147,12 +170,12 @@ fits the work it does:
147
170
 
148
171
  Both serve apps are tuned rather than defaulted:
149
172
 
150
- - **`HuggingFaceImageDeployment`** compiles the denoiser (`transformer` or `unet`)
173
+ - **`Text2Image`** compiles the denoiser (`transformer` or `unet`)
151
174
  and every `text_encoder*` as `Mode.GRAPHED`. A denoising schedule is the ideal case —
152
175
  every step is the same shapes — so the capture is made once and replayed for the
153
176
  whole loop. The VAE is left alone: smallest win of the three, and the one whose
154
177
  shapes move most, especially with tiling on.
155
- - **`HuggingFaceCompletingDeployment`** does *not* compile the module. Decode has
178
+ - **`Text2Text`** does *not* compile the module. Decode has
156
179
  no stable shape of its own — the sequence grows a token per forward — so it
157
180
  pins one: `cache_implementation="static"` with a fixed `max_cache_len`, plus a
158
181
  `CompileConfig(mode=Mode.GRAPHED)`. transformers then compiles `generate`
@@ -196,7 +219,7 @@ it is the floor; only quantizing the weights or running fewer forwards moves it.
196
219
  `complete` streams `CompletionChunk`s (`.content`, `.tool_calls`, `.finish_reason`):
197
220
 
198
221
  ```python
199
- imp = mg.HuggingFaceCompletingImport("Qwen/Qwen2.5-7B-Instruct")
222
+ imp = mg.HuggingFaceImporter("Qwen/Qwen2.5-7B-Instruct", mg.Text2Text)
200
223
  model = imp.client(deployment.url)
201
224
  messages = [{"role": "user", "content": "Explain RAG in one sentence."}]
202
225
 
@@ -211,17 +234,17 @@ async for chunk in mg.complete(model, messages, max_new_tokens=512, temperature=
211
234
  ### Anthropic
212
235
 
213
236
  An Anthropic model is a registry entry like any other — the difference is that
214
- there are no weights to stage, so it is imported with no source, and the serve
215
- app forwards to the API instead of loading anything:
237
+ there are no weights to stage, so there is no importer: the serve app forwards to
238
+ the API instead of loading anything, and a `Hosted` entry registers it:
216
239
 
217
240
  ```python
218
241
  cortexgrid.set_secret("ANTHROPIC_API_KEY", "sk-ant-...") # once
219
242
 
220
- imp = mg.AnthropicImport("claude-sonnet-5")
243
+ entry = mg.Hosted("claude-sonnet-5", mg.AnthropicText2Text) # api_key_secret="..." to use another secret
221
244
  cortexgrid.register_model(
222
- imp.serve_app,
223
- family=imp.family, suffix=imp.suffix,
224
- requirements=imp.requirements(), config=imp.config(),
245
+ entry.serve_app,
246
+ family=entry.family, suffix=entry.suffix,
247
+ requirements=entry.requirements(), config=entry.config(),
225
248
  )
226
249
  ```
227
250
 
@@ -238,45 +261,95 @@ a re-deploy rather than a re-registration.
238
261
 
239
262
  A replica asks for no hardware at all, so it is placed on any node, CPU-only
240
263
  included. From there it deploys and streams exactly like a cluster-served model:
241
- `imp.client(deployment.url)` returns the same `ServedCompletingModel`. Anthropic
264
+ `entry.client(deployment.url)` returns the same `ServedCompletingModel`. Anthropic
242
265
  reports tool calls as structured blocks rather than as generated text, so the
243
266
  serve app re-encodes them into the text form the client parses.
244
267
 
268
+ ### Gemini
269
+
270
+ A Gemini model registers exactly like an Anthropic one, with the serve app for
271
+ its task — `GeminiText2Text` or `GeminiText2Image`. The model id passed to
272
+ `Hosted` is the Gemini model name the deployment calls:
273
+
274
+ ```python
275
+ cortexgrid.set_secret("GEMINI_API_KEY", "AIza...") # once
276
+
277
+ entry = mg.Hosted("gemini-2.5-pro", mg.GeminiText2Text) # or "gemini-2.5-flash", ...
278
+ cortexgrid.register_model(
279
+ entry.serve_app,
280
+ family=entry.family, suffix=entry.suffix,
281
+ requirements=entry.requirements(), config=entry.config(),
282
+ )
283
+ ```
284
+
285
+ Everything said above for Anthropic holds: `config` carries the model name and
286
+ the name of the secret (`api_key_secret="..."` to use another), both editable on
287
+ the model card; a replica needs no hardware; and `entry.client(deployment.url)`
288
+ is a `ServedCompletingModel`, with Gemini's function calls re-encoded into the
289
+ text form the client parses.
290
+
291
+ `GeminiText2Image` answers the same `/generate` route as `Text2Image`, so its
292
+ client is a `ServedGeneratingModel`:
293
+
294
+ ```python
295
+ entry = mg.Hosted("gemini-2.5-flash-image", mg.GeminiText2Image)
296
+ ...
297
+ picture = await mg.generate(entry.client(deployment.url), "a fox in the snow", size=2048)
298
+ ```
299
+
300
+ It takes a prompt only, and makes a square picture. `size` picks Gemini's size
301
+ tier (up to 1K, 2K or 4K) and `seed` is passed on; `steps` and `guidance` are
302
+ diffusion settings with no Gemini counterpart, and come back as `None`.
303
+
245
304
  ## API
246
305
 
247
306
  | | |
248
307
  |---|---|
249
- | `HuggingFaceCompletingImport(hf_id, token=None)` | Importer for a causal LM. |
250
- | `HuggingFaceImageImport(hf_id, token=None, ignore_patterns=None)` | Importer for a diffusers pipeline. |
251
- | `AnthropicImport(model_id, api_key_secret="ANTHROPIC_API_KEY")` | Importer for an Anthropic model; no weights. |
252
- | `ModelImport` | Base of all three: `family`, `suffix`, `serve_app`, `requirements()`, `config()`, `client(url)`. |
253
- | `HuggingFaceImport` | Adds the download half — `source` and the scratch directory. Subclass it for another HuggingFace family. |
308
+ | `HuggingFaceImporter(hf_id, serve_app, token=None, ignore_patterns=None)` | Importer for a model on the HuggingFace Hub, run by `serve_app`. |
309
+ | `Importer` | Base of the importers: identity, scratch directory, `source`, and `requirements()` / `client(url)` taken from the serve app. A new source implements `download(local_dir)` and `weights()`. |
310
+ | `Hosted(model_id, serve_app, **settings)` | Entry for a model hosted elsewhere; no weights. `settings` are the serve app's `config` arguments. |
311
+ | `ModelEntry` | Base of both: `model_id`, `family`, `suffix`, `serve_app`, `requirements()`, `config()`, `client(url)`. |
312
+ | `Text2Text`, `Text2Image`, `Image2Mesh` | Serve apps that run weights (`LocalModel`s). |
313
+ | `AnthropicText2Text` | Serve app that forwards to the Anthropic API (a `HostedModel`). |
314
+ | `GeminiText2Text`, `GeminiText2Image` | Serve apps that forward to the Gemini API (`HostedModel`s sharing the `GeminiModel` base). |
315
+ | `LocalModel` | Base of the serve apps that run weights: `ignore_patterns()`, `bytes_per_param`, `min_vram_gb`, `requirements(weights)`, `client(url, name)`. |
316
+ | `HostedModel` | Base of the serve apps that forward: `config(model_id, **settings)`, `requirements()`, `client(url, name)`. |
317
+ | `Weights(params=None, file_bytes=0)` | What an importer reads about the weights, handed to the serve app to size a replica. |
254
318
  | `complete(model, messages, tools=None, max_new_tokens=2048, temperature=0.7, **kw)` | Async stream of `CompletionChunk` for a `CompletingModel`. |
255
319
  | `generate(model, prompt, *, image=None, **kw) -> GeneratedImage` | One image from a `GeneratingModel`. |
256
320
  | `ServedCompletingModel(url, model_id)` | Client for any deployed completion app. |
257
- | `HuggingFaceImageModel(url, model_id)` | Client for a deployed diffusers app. |
321
+ | `ServedGeneratingModel(url, model_id)` | Client for any deployed text-to-image app. |
322
+ | `ServedMeshingModel(url, model_id)` | Client for any deployed image-to-mesh app. |
258
323
  | `split_model_id(model_id) -> (family, suffix)` | The registry identity an importer derives. |
259
324
  | `detect_device()` | The torch device a serve app should load onto. |
260
325
  | `compiling.Mode.GRAPHED` / `compiling.Mode.FUSED` | The two modes a serve app chooses between. See [Compilation](#compilation). |
261
326
  | `compiling.supported(device)` / `compile_module(module, device, mode)` / `compile_pipeline(pipe, device, mode)` | Whether to compile here, and what a serve app calls to do it. |
262
327
  | `ModelDeployFailed` | Re-exported from cortexgrid; subclasses `RuntimeError`. |
263
328
 
264
- Types: `DeployedModel`, `CompletingModel`, `GeneratingModel`, `CompletionChunk`,
265
- `GeneratedImage`, `ToolCall`, `Message`, `Tool`, `ToolSpec`.
329
+ Types: `DeployedModel`, `CompletingModel`, `GeneratingModel`, `MeshingModel`,
330
+ `CompletionChunk`, `GeneratedImage`, `GeneratedMesh`, `ToolCall`, `Message`, `Tool`,
331
+ `ToolSpec`.
332
+
333
+ Extending it:
266
334
 
267
- Adding a provider means one importer and one serve app. A completion app only has
268
- to stream text from `POST /complete`, inlining tool calls as
269
- `<tool_call>{"name": ..., "arguments": {...}}</tool_call>` (there is an
270
- `encode_tool_call` for upstreams that report them structurally) — then
271
- `ServedCompletingModel` is its client, unchanged.
335
+ - **A new source** is one importer: subclass `Importer` with `download(local_dir)`
336
+ and `weights()`, and every serve app whose weight format it delivers runs its
337
+ models unchanged.
338
+ - **A new hosted API** is one serve app: subclass `HostedModel`, and register it
339
+ with `Hosted`. A text-to-text one only has to stream text from `POST /complete`,
340
+ inlining tool calls as `<tool_call>{"name": ..., "arguments": {...}}</tool_call>`
341
+ (there is an `encode_tool_call` for upstreams that report them structurally) —
342
+ then `ServedCompletingModel` is its client, unchanged.
343
+ - **A model no task's serve app can load** gets its own serve app — an `Image2Mesh`
344
+ subclass, say — and, if its source is unusual too, its own importer.
272
345
 
273
346
  ## Notes
274
347
 
275
348
  - **HuggingFace auth.** For gated/private repos, pass `token=` to the importer (or
276
349
  read `HF_TOKEN` yourself and pass it); it travels with the importer to the node
277
350
  that downloads.
278
- - **Image weight pruning.** `HuggingFaceImageImport` skips example images, docs and
279
- `.gitattributes` by default — nothing `from_pretrained` reads. It does **not**
351
+ - **Image weight pruning.** `Text2Image` tells the importer to skip example images,
352
+ docs and `.gitattributes` — nothing `from_pretrained` reads. It does **not**
280
353
  skip a repo's consolidated single-file checkpoint, which is weights it cannot
281
354
  tell apart from the component ones: name it per-model in
282
355
  `HF_IMAGE_SNAPSHOT_IGNORE` (comma-separated globs, e.g.
@@ -286,7 +359,7 @@ to stream text from `POST /complete`, inlining tool calls as
286
359
  describe what actually gets staged:
287
360
 
288
361
  ```python
289
- >>> mg.HuggingFaceImageImport("black-forest-labs/FLUX.2-klein-base-4B").requirements()
362
+ >>> mg.HuggingFaceImporter("black-forest-labs/FLUX.2-klein-base-4B", mg.Text2Image).requirements()
290
363
  ModelRequirements(num_gpus=1, ram_gb=26.3, vram_gb=29.6) # single-file checkpoint included
291
364
  ```
292
365
  - **Only model files are stored.** HuggingFace's download bookkeeping