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

{cortexgrid_infer-0.3.1 → cortexgrid_infer-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid-infer
-Version: 0.3.1
+Version: 0.3.2
 Summary: Serve apps per task (text-to-text, text-to-image, image-to-mesh) and importers per source (HuggingFace) for running models on cortexgrid
 Project-URL: Homepage, https://github.com/robodatalab/cortexgrid-infer
 Project-URL: Repository, https://github.com/robodatalab/cortexgrid-infer
@@ -12,6 +12,7 @@ Requires-Dist: anthropic>=0.40.0
 Requires-Dist: cortexgrid>=0.3.4
 Requires-Dist: diffusers>=0.38
 Requires-Dist: fastapi>=0.133
+Requires-Dist: google-genai>=1.0
 Requires-Dist: httpx>=0.27.0
 Requires-Dist: huggingface-hub>=0.20.0
 Requires-Dist: numpy>=1.24
@@ -63,6 +64,8 @@ A model is two things, and the library keeps them apart:
 | `Text2Image` | text (and image) to image | any diffusers pipeline | `ServedGeneratingModel` (`generate`) |
 | `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) |
 | `AnthropicText2Text` | text to text | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
+| `GeminiText2Text` | text to text | the Gemini API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
+| `GeminiText2Image` | text to image | the Gemini API, forwarded from the cluster | `ServedGeneratingModel` (`generate`) |
 
 | Importer | Source |
 |---|---|
@@ -290,6 +293,42 @@ included. From there it deploys and streams exactly like a cluster-served model:
 reports tool calls as structured blocks rather than as generated text, so the
 serve app re-encodes them into the text form the client parses.
 
+### Gemini
+
+A Gemini model registers exactly like an Anthropic one, with the serve app for
+its task — `GeminiText2Text` or `GeminiText2Image`. The model id passed to
+`Hosted` is the Gemini model name the deployment calls:
+
+```python
+cortexgrid.set_secret("GEMINI_API_KEY", "AIza...")   # once
+
+entry = mg.Hosted("gemini-2.5-pro", mg.GeminiText2Text)   # or "gemini-2.5-flash", ...
+cortexgrid.register_model(
+    entry.serve_app,
+    family=entry.family, suffix=entry.suffix,
+    requirements=entry.requirements(), config=entry.config(),
+)
+```
+
+Everything said above for Anthropic holds: `config` carries the model name and
+the name of the secret (`api_key_secret="..."` to use another), both editable on
+the model card; a replica needs no hardware; and `entry.client(deployment.url)`
+is a `ServedCompletingModel`, with Gemini's function calls re-encoded into the
+text form the client parses.
+
+`GeminiText2Image` answers the same `/generate` route as `Text2Image`, so its
+client is a `ServedGeneratingModel`:
+
+```python
+entry = mg.Hosted("gemini-2.5-flash-image", mg.GeminiText2Image)
+...
+picture = await mg.generate(entry.client(deployment.url), "a fox in the snow", size=2048)
+```
+
+It takes a prompt only, and makes a square picture. `size` picks Gemini's size
+tier (up to 1K, 2K or 4K) and `seed` is passed on; `steps` and `guidance` are
+diffusion settings with no Gemini counterpart, and come back as `None`.
+
 ## API
 
 | | |
@@ -300,6 +339,7 @@ serve app re-encodes them into the text form the client parses.
 | `ModelEntry` | Base of both: `model_id`, `family`, `suffix`, `serve_app`, `requirements()`, `config()`, `client(url)`. |
 | `Text2Text`, `Text2Image`, `Image2Mesh` | Serve apps that run weights (`LocalModel`s). |
 | `AnthropicText2Text` | Serve app that forwards to the Anthropic API (a `HostedModel`). |
+| `GeminiText2Text`, `GeminiText2Image` | Serve apps that forward to the Gemini API (`HostedModel`s sharing the `GeminiModel` base). |
 | `LocalModel` | Base of the serve apps that run weights: `ignore_patterns()`, `bytes_per_param`, `min_vram_gb`, `requirements(weights)`, `client(url, name)`. |
 | `HostedModel` | Base of the serve apps that forward: `config(model_id, **settings)`, `requirements()`, `client(url, name)`. |
 | `Weights(params=None, file_bytes=0)` | What an importer reads about the weights, handed to the serve app to size a replica. |

{cortexgrid_infer-0.3.1 → cortexgrid_infer-0.3.2}/README.md

@@ -36,6 +36,8 @@ A model is two things, and the library keeps them apart:
 | `Text2Image` | text (and image) to image | any diffusers pipeline | `ServedGeneratingModel` (`generate`) |
 | `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) |
 | `AnthropicText2Text` | text to text | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
+| `GeminiText2Text` | text to text | the Gemini API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
+| `GeminiText2Image` | text to image | the Gemini API, forwarded from the cluster | `ServedGeneratingModel` (`generate`) |
 
 | Importer | Source |
 |---|---|
@@ -263,6 +265,42 @@ included. From there it deploys and streams exactly like a cluster-served model:
 reports tool calls as structured blocks rather than as generated text, so the
 serve app re-encodes them into the text form the client parses.
 
+### Gemini
+
+A Gemini model registers exactly like an Anthropic one, with the serve app for
+its task — `GeminiText2Text` or `GeminiText2Image`. The model id passed to
+`Hosted` is the Gemini model name the deployment calls:
+
+```python
+cortexgrid.set_secret("GEMINI_API_KEY", "AIza...")   # once
+
+entry = mg.Hosted("gemini-2.5-pro", mg.GeminiText2Text)   # or "gemini-2.5-flash", ...
+cortexgrid.register_model(
+    entry.serve_app,
+    family=entry.family, suffix=entry.suffix,
+    requirements=entry.requirements(), config=entry.config(),
+)
+```
+
+Everything said above for Anthropic holds: `config` carries the model name and
+the name of the secret (`api_key_secret="..."` to use another), both editable on
+the model card; a replica needs no hardware; and `entry.client(deployment.url)`
+is a `ServedCompletingModel`, with Gemini's function calls re-encoded into the
+text form the client parses.
+
+`GeminiText2Image` answers the same `/generate` route as `Text2Image`, so its
+client is a `ServedGeneratingModel`:
+
+```python
+entry = mg.Hosted("gemini-2.5-flash-image", mg.GeminiText2Image)
+...
+picture = await mg.generate(entry.client(deployment.url), "a fox in the snow", size=2048)
+```
+
+It takes a prompt only, and makes a square picture. `size` picks Gemini's size
+tier (up to 1K, 2K or 4K) and `seed` is passed on; `steps` and `guidance` are
+diffusion settings with no Gemini counterpart, and come back as `None`.
+
 ## API
 
 | | |
@@ -273,6 +311,7 @@ serve app re-encodes them into the text form the client parses.
 | `ModelEntry` | Base of both: `model_id`, `family`, `suffix`, `serve_app`, `requirements()`, `config()`, `client(url)`. |
 | `Text2Text`, `Text2Image`, `Image2Mesh` | Serve apps that run weights (`LocalModel`s). |
 | `AnthropicText2Text` | Serve app that forwards to the Anthropic API (a `HostedModel`). |
+| `GeminiText2Text`, `GeminiText2Image` | Serve apps that forward to the Gemini API (`HostedModel`s sharing the `GeminiModel` base). |
 | `LocalModel` | Base of the serve apps that run weights: `ignore_patterns()`, `bytes_per_param`, `min_vram_gb`, `requirements(weights)`, `client(url, name)`. |
 | `HostedModel` | Base of the serve apps that forward: `config(model_id, **settings)`, `requirements()`, `client(url, name)`. |
 | `Weights(params=None, file_bytes=0)` | What an importer reads about the weights, handed to the serve app to size a replica. |

{cortexgrid_infer-0.3.1 → cortexgrid_infer-0.3.2}/cortexgrid_infer/__init__.py

@@ -50,6 +50,8 @@ from cortexgrid_infer.meshing import ServedMeshingModel
 from cortexgrid_infer.device import detect_device
 from cortexgrid_infer.models import (
     AnthropicText2Text,
+    GeminiText2Image,
+    GeminiText2Text,
     HostedModel,
     Image2Mesh,
     LocalModel,
@@ -87,6 +89,8 @@ __all__ = [
     "Text2Image",
     "Image2Mesh",
     "AnthropicText2Text",
+    "GeminiText2Text",
+    "GeminiText2Image",
     "ModelEntry",
     "Hosted",
     "split_model_id",

{cortexgrid_infer-0.3.1 → cortexgrid_infer-0.3.2}/cortexgrid_infer/models/__init__.py

@@ -3,6 +3,7 @@ that task whichever source it came from."""
 
 from cortexgrid_infer.models.anthropic import AnthropicText2Text
 from cortexgrid_infer.models.base import HostedModel, LocalModel, Weights
+from cortexgrid_infer.models.gemini import GeminiText2Image, GeminiText2Text
 from cortexgrid_infer.models.image2mesh import Image2Mesh
 from cortexgrid_infer.models.text2image import Text2Image
 from cortexgrid_infer.models.text2text import Text2Text
@@ -15,4 +16,6 @@ __all__ = [
     "Text2Image",
     "Image2Mesh",
     "AnthropicText2Text",
+    "GeminiText2Text",
+    "GeminiText2Image",
 ]

cortexgrid_infer-0.3.2/cortexgrid_infer/models/gemini/__init__.py

@@ -0,0 +1,11 @@
+"""The serve apps that forward to the Gemini API, one per task."""
+
+from cortexgrid_infer.models.gemini.base import GeminiModel
+from cortexgrid_infer.models.gemini.text2image import GeminiText2Image
+from cortexgrid_infer.models.gemini.text2text import GeminiText2Text
+
+__all__ = [
+    "GeminiModel",
+    "GeminiText2Text",
+    "GeminiText2Image",
+]

cortexgrid_infer-0.3.2/cortexgrid_infer/models/gemini/base.py

@@ -0,0 +1,50 @@
+"""What every Gemini serve app shares: which model to call, and the key to call
+it with.
+
+A Gemini app loads no weights and needs no GPU, so there is nothing to import:
+it is registered with `registry.Hosted`. What it needs instead comes from the
+registry entry's `config`, set when the model was registered and editable on its
+model card afterwards.
+"""
+
+from __future__ import annotations
+
+import cortexgrid
+from google import genai
+
+from cortexgrid_infer.models.base import HostedModel
+
+# Keys the registry entry must carry for the deployment to reach the API.
+MODEL_PARAM = "model"
+API_KEY_SECRET_PARAM = "api_key_secret"
+
+# The cortexgrid secret the deployment reads its API key from. It is the secret's
+# name that is stored on the registry entry, never the key itself: entries are
+# readable by anyone who can see the model.
+DEFAULT_API_KEY_SECRET = "GEMINI_API_KEY"
+
+
+class GeminiModel(HostedModel):
+    """Base of the serve apps that forward to the Gemini API, one per task."""
+
+    @classmethod
+    def config(
+        cls, model_id: str, api_key_secret: str = DEFAULT_API_KEY_SECRET
+    ) -> dict[str, str]:
+        """`model_id` is the name Gemini knows the model by, and is what the
+        deployment sends upstream; `api_key_secret` names the cortexgrid secret
+        holding the key to send it with."""
+        return {MODEL_PARAM: model_id, API_KEY_SECRET_PARAM: api_key_secret}
+
+    def __init__(self, family: str, suffix: str, run_name: str) -> None:
+        config = cortexgrid.model_config(family, suffix, run_name)
+        missing = {MODEL_PARAM, API_KEY_SECRET_PARAM} - config.keys()
+        if missing:
+            raise RuntimeError(
+                f"{family}/{suffix}/{run_name} is missing {sorted(missing)} from "
+                "its config; set them on the model card"
+            )
+        self._model = config[MODEL_PARAM]
+        self._client = genai.Client(
+            api_key=cortexgrid.get_secret(config[API_KEY_SECRET_PARAM])
+        )

cortexgrid_infer-0.3.2/cortexgrid_infer/models/gemini/text2image.py

@@ -0,0 +1,116 @@
+"""The text-to-image serve app that forwards to the Gemini API.
+
+It answers the same `POST /generate` route as `Text2Image` (see
+`cortexgrid_infer.imaging`), so the same client talks to it. It takes a prompt
+only. Of the sampling options, `size` and `seed` have a Gemini counterpart;
+`steps` and `guidance` are diffusion settings, and are reported back as unused.
+"""
+
+from __future__ import annotations
+
+import base64
+import io
+import time
+from typing import Any
+
+from cortexgrid import serve
+from fastapi import FastAPI, HTTPException
+from google.genai import types
+from PIL import Image
+
+from cortexgrid_infer.imaging import ServedGeneratingModel
+from cortexgrid_infer.models.gemini.base import GeminiModel
+
+_app = FastAPI()
+
+
+def image_size(size: int) -> str | None:
+    """The Gemini size tier that covers `size` pixels a side.
+
+    None up to 1K, the model's default, so that it is left unsent: not every
+    Gemini image model takes a size tier."""
+    if size <= 1024:
+        return None
+    if size <= 2048:
+        return "2K"
+    return "4K"
+
+
+def to_png(data: bytes) -> tuple[bytes, int, int]:
+    """`data`, a picture in whatever format Gemini returned it, as PNG bytes
+    along with its width and height."""
+    image = Image.open(io.BytesIO(data))
+    if image.format == "PNG":
+        return data, image.width, image.height
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    return buffer.getvalue(), image.width, image.height
+
+
+def no_image_reason(response: types.GenerateContentResponse) -> str:
+    """Why `response` carries no picture: what the model said instead, or why
+    it stopped."""
+    feedback = response.prompt_feedback
+    if feedback and feedback.block_reason:
+        return f"the prompt was blocked ({feedback.block_reason})"
+    text = " ".join(part.text for part in response.parts or [] if part.text)
+    if text:
+        return text
+    candidate = response.candidates[0] if response.candidates else None
+    return f"finish reason {candidate.finish_reason if candidate else None}"
+
+
+@serve.ingress(_app)
+class GeminiText2Image(GeminiModel):
+    @classmethod
+    def client(cls, url: str, name: str) -> ServedGeneratingModel:
+        return ServedGeneratingModel(url=url, model_id=name)
+
+    @_app.post("/generate")
+    async def generate(self, body: dict[str, Any]) -> dict[str, Any]:
+        if body.get("image"):
+            raise HTTPException(
+                status_code=400,
+                detail=f"{type(self).__name__} takes a prompt only, not an image",
+            )
+        size = int(body.get("size") or 1024)
+        seed = body.get("seed")
+
+        # Square, like `Text2Image` without a reference picture.
+        image_config: dict[str, Any] = {"aspect_ratio": "1:1"}
+        tier = image_size(size)
+        if tier:
+            image_config["image_size"] = tier
+        config: dict[str, Any] = {
+            "response_modalities": ["IMAGE"],
+            "image_config": image_config,
+        }
+        if seed is not None:
+            config["seed"] = int(seed)
+
+        t0 = time.time()
+        response = await self._client.aio.models.generate_content(
+            model=self._model, contents=body["prompt"], config=config
+        )
+        duration = time.time() - t0
+
+        data = next(
+            (part.inline_data.data for part in response.parts or [] if part.inline_data),
+            None,
+        )
+        if data is None:
+            raise HTTPException(
+                status_code=502,
+                detail=f"{self._model} returned no image: {no_image_reason(response)}",
+            )
+
+        png, width, height = to_png(data)
+        return {
+            "image": base64.b64encode(png).decode("ascii"),
+            "width": width,
+            "height": height,
+            "steps": None,
+            "guidance": None,
+            "seed": seed,
+            "duration_s": round(duration, 2),
+        }

cortexgrid_infer-0.3.2/cortexgrid_infer/models/gemini/text2text.py

@@ -0,0 +1,146 @@
+"""The text-to-text serve app that forwards to the Gemini API.
+
+It speaks the same `/complete` protocol as `Text2Text`, so the same client talks
+to it. Gemini reports tool calls as structured parts rather than as generated
+text, so they are re-encoded on the way out.
+"""
+
+from __future__ import annotations
+
+import base64
+from collections.abc import AsyncIterator
+from typing import Any
+
+from cortexgrid import serve
+from fastapi import FastAPI
+from fastapi.responses import StreamingResponse
+
+from cortexgrid_infer.completion import ServedCompletingModel, encode_tool_call
+from cortexgrid_infer.core import Message, ToolSpec
+from cortexgrid_infer.models.gemini.base import GeminiModel
+
+_app = FastAPI()
+
+# Gemini 3 rejects a function call replayed without the thought signature it was
+# returned with, and signatures do not survive the trip through the client as
+# text. Google documents this value as the one that skips the check. The SDK
+# base64-encodes the bytes it is given, so they are decoded here for the value
+# to go out verbatim.
+SKIP_THOUGHT_SIGNATURE = base64.urlsafe_b64decode("skip_thought_signature_validator")
+
+
+def to_gemini_contents(
+    messages: list[Message],
+) -> tuple[str | None, list[dict[str, Any]]]:
+    """Convert OpenAI-style messages to Gemini contents."""
+    system_prompt: str | None = None
+    contents: list[dict[str, Any]] = []
+    # Gemini answers a function call by name, which a tool message does not
+    # carry, so it is looked up from the call the message answers.
+    tool_names: dict[str, str] = {}
+
+    for msg in messages:
+        role = msg["role"]
+
+        if role == "system":
+            system_prompt = msg["content"]
+
+        elif role == "user":
+            contents.append({"role": "user", "parts": [{"text": msg["content"]}]})
+
+        elif role == "assistant":
+            parts: list[dict[str, Any]] = []
+            if msg.get("content"):
+                parts.append({"text": msg["content"]})
+            for i, tc in enumerate(msg.get("tool_calls", [])):
+                func = tc["function"]
+                tool_names[tc["id"]] = func["name"]
+                call_part: dict[str, Any] = {
+                    "function_call": {
+                        "id": tc["id"],
+                        "name": func["name"],
+                        "args": func["arguments"],
+                    }
+                }
+                # Gemini signs only the first call of a parallel batch.
+                if i == 0:
+                    call_part["thought_signature"] = SKIP_THOUGHT_SIGNATURE
+                parts.append(call_part)
+            contents.append({"role": "model", "parts": parts})
+
+        elif role == "tool":
+            response_part = {
+                "function_response": {
+                    "id": msg["tool_call_id"],
+                    "name": tool_names[msg["tool_call_id"]],
+                    "response": {"output": msg["content"]},
+                }
+            }
+            if (
+                contents
+                and contents[-1]["role"] == "user"
+                and "function_response" in contents[-1]["parts"][0]
+            ):
+                contents[-1]["parts"].append(response_part)
+            else:
+                contents.append({"role": "user", "parts": [response_part]})
+
+    return system_prompt, contents
+
+
+def to_gemini_tools(
+    tool_specs: list[ToolSpec] | None,
+) -> list[dict[str, Any]] | None:
+    """Convert OpenAI-style tool specs to one Gemini tool of function
+    declarations."""
+    if not tool_specs:
+        return None
+    declarations = []
+    for spec in tool_specs:
+        func = spec["function"]
+        declaration = {"name": func["name"], "description": func.get("description", "")}
+        # Plain JSON Schema, which `parameters` (an OpenAPI subset) is not. A
+        # function that takes nothing declares no schema at all.
+        if "parameters" in func:
+            declaration["parameters_json_schema"] = func["parameters"]
+        declarations.append(declaration)
+    return [{"function_declarations": declarations}]
+
+
+@serve.ingress(_app)
+class GeminiText2Text(GeminiModel):
+    @classmethod
+    def client(cls, url: str, name: str) -> ServedCompletingModel:
+        return ServedCompletingModel(url=url, model_id=name)
+
+    @_app.post("/complete")
+    async def complete(self, body: dict[str, Any]) -> StreamingResponse:
+        system_prompt, contents = to_gemini_contents(body["messages"])
+        tools = to_gemini_tools(body.get("tools"))
+
+        generate_config: dict[str, Any] = {
+            "max_output_tokens": body.get("max_new_tokens", 16 * 1024),
+            "temperature": body.get("temperature", 0.7),
+        }
+        if system_prompt:
+            generate_config["system_instruction"] = system_prompt
+        if tools:
+            generate_config["tools"] = tools
+
+        async def stream() -> AsyncIterator[bytes]:
+            chunks = await self._client.aio.models.generate_content_stream(
+                model=self._model, contents=contents, config=generate_config
+            )
+            async for chunk in chunks:
+                if not chunk.candidates or not chunk.candidates[0].content:
+                    continue
+                for part in chunk.candidates[0].content.parts or []:
+                    if part.function_call:
+                        yield encode_tool_call(
+                            part.function_call.name, dict(part.function_call.args or {})
+                        ).encode("utf-8")
+
+                    elif part.text and not part.thought:
+                        yield part.text.encode("utf-8")
+
+        return StreamingResponse(stream(), media_type="text/plain")

{cortexgrid_infer-0.3.1 → cortexgrid_infer-0.3.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cortexgrid-infer"
-version = "0.3.1"
+version = "0.3.2"
 description = "Serve apps per task (text-to-text, text-to-image, image-to-mesh) and importers per source (HuggingFace) for running models on cortexgrid"
 readme = "README.md"
 license = "Apache-2.0"
@@ -22,6 +22,9 @@ dependencies = [
     "transformers>=4.47.0",
     "huggingface_hub>=0.20.0",
     "anthropic>=0.40.0",
+    # models.gemini. The unified Google Gen AI SDK, not the deprecated
+    # google-generativeai.
+    "google-genai>=1.0",
     "python-dotenv>=1.0.0",
     "httpx>=0.27.0",
     # The serve-apps run on the cluster's FastAPI, which Ray 2.58's Serve