Skip to content

image gen aspect ratio #161

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
adambalogh merged 2 commits into from
Sep 9, 2026
Merged

image gen aspect ratio #161

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
feeacf7
image aspect ratio
dixitaniket Sep 9, 2026
6ccd714
fixes
dixitaniket Sep 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions tee_gateway/controllers/chat_controller.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,10 @@
canonical_user_content,
)
from tee_gateway.image_generation import (
aspect_ratio_params,
create_image_generation_response,
create_image_generation_streaming_response,
validate_aspect_ratio,
)
from tee_gateway.model_registry import get_model_config
from tee_gateway.moderation import (
Expand Down Expand Up @@ -102,6 +104,13 @@ def create_chat_completion(body):
except AttachmentValidationError as e:
return {"error": "Invalid attachment", "message": str(e)}, 400

# A ratio the target model can't produce is a client error, so it is caught
# here rather than surfacing from the provider call as a 500.
try:
validate_aspect_ratio(chat_request.model, chat_request.aspect_ratio)
except ValueError as e:
return {"error": "Invalid aspect_ratio", "message": str(e)}, 400

# Score the newest user turn of image requests before any provider work
# (text chat is not moderated — see moderation.should_moderate_model).
# Fail-open: an unavailable moderation endpoint yields an unchecked
Expand Down Expand Up @@ -299,6 +308,16 @@ def _create_non_streaming_response(
if tools_list:
model = model.bind_tools(tools_list)

# Aspect ratio for Gemini's inline-image models rides in Gemini's
# image_config (these models leave image_aspect_ratio_param at its
# default, which is ImageConfig's own field name). Bound after
# bind_tools, which re-binds the base model and drops kwargs bound
# before it. Empty params mean the automatic path: bind nothing.
if cfg.image_output:
image_config = aspect_ratio_params(cfg, chat_request.aspect_ratio)
if image_config:
model = model.bind(image_config=image_config)

# Bind response_format if provided (json_object or json_schema).
# Anthropic does not support response_format via bind(); use
# with_structured_output() for json_schema instead (json_object has no
Expand Down Expand Up @@ -477,6 +496,13 @@ def _create_streaming_response(
if tools_list:
model = model.bind_tools(tools_list)

# Aspect ratio for the inline-image models (see the non-streaming path
# for why this is bound after bind_tools).
if image_output_model:
image_config = aspect_ratio_params(cfg, chat_request.aspect_ratio)
if image_config:
model = model.bind(image_config=image_config)

# Bind response_format if provided (json_object or json_schema).
# Anthropic does not support response_format via bind(); use
# with_structured_output() for json_schema instead (json_object has no
Expand Down Expand Up @@ -956,6 +982,8 @@ def _chat_request_to_dict(chat_request: CreateChatCompletionRequest) -> dict:
d["response_format"] = _normalize_response_format(chat_request.response_format)
if chat_request.web_search:
d["web_search"] = True
if chat_request.aspect_ratio:
d["aspect_ratio"] = chat_request.aspect_ratio
return d


Expand All @@ -979,6 +1007,7 @@ def _parse_chat_request(chat_request_dict: dict) -> CreateChatCompletionRequest:
tool_choice=chat_request_dict.get("tool_choice"),
user=chat_request_dict.get("user"),
web_search=chat_request_dict.get("web_search", False),
aspect_ratio=chat_request_dict.get("aspect_ratio"),
)


Expand Down
56 changes: 54 additions & 2 deletions tee_gateway/image_generation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -290,7 +290,11 @@ def _build_reference_uploads(


def _build_generations_payload(
cfg: Any, prompt: str, count: int, refs: Optional[List[str]]
cfg: Any,
prompt: str,
count: int,
refs: Optional[List[str]],
aspect_ratio: Optional[str] = None,
) -> dict[str, Any]:
"""Build the JSON body for a ``/images/generations`` request.

Expand All @@ -305,16 +309,60 @@ def _build_generations_payload(
payload["n"] = count
if cfg.image_extra_params:
payload.update(cfg.image_extra_params)
payload.update(aspect_ratio_params(cfg, aspect_ratio))
if refs:
payload["image"] = refs[0] if len(refs) == 1 else refs
return payload


def aspect_ratio_params(cfg: Any, aspect_ratio: Optional[str]) -> dict[str, str]:
"""Translate the public ratio into this provider's request shape.

Returns an empty dict for the automatic path (field omitted, or ``auto``),
so callers must treat "no params" as "let the provider choose" rather than
indexing the result.
"""
if aspect_ratio is None:
return {}
if not isinstance(aspect_ratio, str):
raise ValueError("aspect_ratio must be a string")
ratio = aspect_ratio.strip()
if not ratio or ratio == "auto":
return {}
supported = cfg.image_aspect_ratios or {}
value = supported.get(ratio)
if value is None:
choices = ", ".join(supported) or "none"
raise ValueError(
f"Unsupported aspect_ratio {ratio!r} for this model; supported: {choices}"
)
return {cfg.image_aspect_ratio_param: value}


def validate_aspect_ratio(model: str, aspect_ratio: Optional[str]) -> None:
"""Reject an unsupported ratio before any provider work is done.

Raises ``ValueError`` so the controller can answer 400 rather than letting
the same error surface from deep inside a generation as a 500. Models that
cannot shape their output (every text model) ignore the field instead of
failing the request — like the deprecated ``web_search`` flag.
"""
if aspect_ratio is None:
return
try:
cfg = get_model_config(model)
except ValueError:
return # unknown model: reported by the normal request path
if cfg.image_generation or cfg.image_output:
aspect_ratio_params(cfg, aspect_ratio)


def generate_images(
model: str,
prompt: str,
n: int = 1,
reference_images: Optional[List[str]] = None,
aspect_ratio: Optional[str] = None,
) -> tuple[list[str], int]:
"""Generate images via a provider's OpenAI-compatible images endpoint.

Expand Down Expand Up @@ -382,6 +430,7 @@ def generate_images(
form["response_format"] = cfg.image_response_format
if cfg.image_extra_params:
form.update({k: str(v) for k, v in cfg.image_extra_params.items()})
form.update(aspect_ratio_params(cfg, aspect_ratio))
resp = client.post(edit_endpoint, data=form, files=uploads)
else:
# No edit endpoint (or nothing uploadable): JSON generations. Inline
Expand All @@ -390,7 +439,9 @@ def generate_images(
json_refs = refs if not cfg.image_edit_endpoint else None
resp = client.post(
_IMAGE_GENERATION_PATH,
json=_build_generations_payload(cfg, prompt, count, json_refs),
json=_build_generations_payload(
cfg, prompt, count, json_refs, aspect_ratio=aspect_ratio
),
)
_raise_for_status_with_detail(resp)
data = resp.json().get("data", []) or []
Expand Down Expand Up @@ -497,6 +548,7 @@ def _run_image_generation(
prompt,
n=chat_request.n or 1,
reference_images=reference_images,
aspect_ratio=chat_request.aspect_ratio,
)

timestamp = int(time.time())
Expand Down
134 changes: 129 additions & 5 deletions tee_gateway/model_registry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,16 @@ class ModelConfig:
# Static extra params merged verbatim into the request payload (e.g. size,
# watermark). Keyed by field name; values must be JSON-serializable.
image_extra_params: Optional[Mapping[str, Any]] = None
# Optional aspect-ratio choices exposed to callers. Keys are normalized
# public ``aspect_ratio`` values and values are provider-specific values.
# Omitted/"auto" requests send no ratio override. Required on any image
# model: without it, callers get no shape choice at all.
image_aspect_ratios: Optional[Mapping[str, str]] = None
# Request field the ratio is sent under — a provider taking explicit pixel
# dimensions uses "size". ``image_output`` (inline-image chat) models must
# leave this at the default: it doubles as the field name inside Gemini's
# ImageConfig, which the chat controller binds the ratio through.
image_aspect_ratio_param: str = "aspect_ratio"
# USD per image-modality output token, for ``image_output`` models (Gemini
# "nano banana"). These providers bill image output at a higher rate than
# text/thinking output: image tokens at this rate, text + thinking tokens at
Expand Down Expand Up @@ -117,6 +127,83 @@ class ModelConfig:
"stream": False,
}

# 2K-ish explicit dimensions for ModelArk models. The bare ``size: 2K`` config
# remains the auto path, where ModelArk chooses the shape from the prompt or
# reference image.
_BYTEDANCE_2K_ASPECT_SIZES: dict[str, str] = {
"1:1": "2048x2048",
"16:9": "2560x1440",
"9:16": "1440x2560",
"4:3": "2368x1776",
"3:4": "1776x2368",
"3:2": "2496x1664",
"2:3": "1664x2496",
"21:9": "3136x1344",
}

# gpt-image-2 takes explicit ``WIDTHxHEIGHT`` sizes rather than ratios: it
# accepts any size whose edges are multiples of 16, whose long:short ratio is
# at most 3:1, whose longest edge is <= 3840px, and whose pixel count is
# between 655,360 and 8,294,400. OpenAI's three documented presets
# (1024x1024, 1536x1024, 1024x1536) are kept verbatim; the rest are the
# exact-ratio sizes closest to them in pixel count, so every shape stays
# inside the ~1.0-1.6MP band the flat per-image price is set for (output
# tokens, and so our real cost, scale with pixel count). Omitting the field
# leaves size at OpenAI's ``auto``.
_GPT_IMAGE_ASPECT_SIZES: dict[str, str] = {
"1:1": "1024x1024",
"3:2": "1536x1024",
"2:3": "1024x1536",
"4:3": "1408x1056",
"3:4": "1056x1408",
"5:4": "1280x1024",
"4:5": "1024x1280",
"16:9": "1536x864",
"9:16": "864x1536",
"21:9": "1792x768",
}

_GEMINI_IMAGE_ASPECT_RATIOS: dict[str, str] = {
ratio: ratio
for ratio in (
"1:1",
"1:4",
"1:8",
"2:3",
"3:2",
"3:4",
"4:1",
"4:3",
"4:5",
"5:4",
"8:1",
"9:16",
"16:9",
"21:9",
)
}

_GROK_IMAGE_ASPECT_RATIOS: dict[str, str] = {
ratio: ratio
for ratio in (
"1:1",
"16:9",
"9:16",
"4:3",
"3:4",
"3:2",
"2:3",
"2:1",
"1:2",
"19.5:9",
"9:19.5",
"20:9",
"9:20",
"21:9",
"5:2",
)
}


@unique
class SupportedModel(Enum):
Expand Down Expand Up @@ -240,9 +327,9 @@ class SupportedModel(Enum):
# live on OpenAI's separate ``/images/edits`` endpoint, which takes the
# reference images as multipart file uploads rather than a JSON ``image``
# field — so reference turns are routed there via ``image_edit_endpoint``
# (up to 10 references per request). Size/quality are pinned so the flat
# per-image price stays predictable. Billed at a flat $0.05 per generated
# image; token prices unused.
# (up to 10 references per request). Quality remains pinned to medium while
# size defaults to auto unless the caller selects a shape.
# Billed at a flat $0.05 per generated image; token prices unused.
GPT_IMAGE_2 = ModelConfig(
provider="openai",
api_name="gpt-image-2",
Expand All @@ -253,7 +340,9 @@ class SupportedModel(Enum):
image_response_format=None,
image_supports_reference=True,
image_edit_endpoint="/images/edits",
image_extra_params={"size": "1024x1024", "quality": "medium"},
image_extra_params={"quality": "medium"},
image_aspect_ratios=_GPT_IMAGE_ASPECT_SIZES,
image_aspect_ratio_param="size",
)

# ── Anthropic ───────────────────────────────────────────────────────
Expand Down Expand Up @@ -392,6 +481,21 @@ class SupportedModel(Enum):
output_price_usd=Decimal("0.0000015"),
image_output=True,
image_output_price_usd=Decimal("0.00003"),
image_aspect_ratios={
ratio: ratio
for ratio in (
"1:1",
"2:3",
"3:2",
"3:4",
"4:3",
"4:5",
"5:4",
"9:16",
"16:9",
"21:9",
)
},
)
# Native image generation ("nano banana 2"), the latest Gemini image model.
# Google bills output at two rates: text/thinking at $3/MTok and images at
Expand All @@ -404,6 +508,7 @@ class SupportedModel(Enum):
output_price_usd=Decimal("0.000003"),
image_output=True,
image_output_price_usd=Decimal("0.00006"),
image_aspect_ratios=_GEMINI_IMAGE_ASPECT_RATIOS,
)
GEMINI_3_5_FLASH = ModelConfig(
provider="google",
Expand Down Expand Up @@ -510,6 +615,7 @@ class SupportedModel(Enum):
image_generation=True,
per_image_price_usd=Decimal("0.02"),
image_response_format="url",
image_aspect_ratios=_GROK_IMAGE_ASPECT_RATIOS,
)
# Grok Imagine Image 2.0 — xAI's newer, higher-quality image model
# (released ~2026-08-11), offered alongside grok-imagine-image rather than
Expand All @@ -527,6 +633,7 @@ class SupportedModel(Enum):
image_generation=True,
per_image_price_usd=Decimal("0.04"),
image_response_format="url",
image_aspect_ratios=_GROK_IMAGE_ASPECT_RATIOS,
)

# ── ByteDance (BytePlus ModelArk, OpenAI-compatible) ────────────────
Expand Down Expand Up @@ -579,6 +686,8 @@ class SupportedModel(Enum):
image_generation=True,
per_image_price_usd=Decimal("0.03"),
image_supports_reference=True,
image_aspect_ratios=_BYTEDANCE_2K_ASPECT_SIZES,
image_aspect_ratio_param="size",
)
# Seedream 5.0 Lite image generation via a ModelArk deployment endpoint.
# Seedream 5.0 Lite image generation via a ModelArk deployment endpoint
Expand All @@ -595,6 +704,8 @@ class SupportedModel(Enum):
image_send_n=False,
image_supports_reference=True,
image_extra_params=_BYTEDANCE_EP_IMAGE_PARAMS,
image_aspect_ratios=_BYTEDANCE_2K_ASPECT_SIZES,
image_aspect_ratio_param="size",
)
# Seedance 4.5 image generation via a ModelArk deployment endpoint.
# Returns hosted URLs (fetched and inlined by the gateway) and takes the
Expand All @@ -610,6 +721,8 @@ class SupportedModel(Enum):
image_send_n=False,
image_supports_reference=True,
image_extra_params=_BYTEDANCE_EP_IMAGE_PARAMS,
image_aspect_ratios=_BYTEDANCE_2K_ASPECT_SIZES,
image_aspect_ratio_param="size",
)
# Seedance 5.0 image generation via a ModelArk deployment endpoint.
# Returns hosted URLs (fetched and inlined by the gateway). Unlike the
Expand All @@ -628,6 +741,8 @@ class SupportedModel(Enum):
image_send_n=False,
image_supports_reference=True,
image_extra_params=_SEEDANCE_5_IMAGE_PARAMS,
image_aspect_ratios=_BYTEDANCE_2K_ASPECT_SIZES,
image_aspect_ratio_param="size",
)

# ── OpenRouter (OpenAI-compatible) ──────────────────────────────────
Expand Down Expand Up @@ -674,7 +789,16 @@ class SupportedModel(Enum):
per_image_price_usd=Decimal("0.015"),
image_response_format=None,
image_send_n=False,
image_extra_params={"size": "1280x1280"},
image_aspect_ratios={
"1:1": "1280x1280",
"3:2": "1568x1056",
"2:3": "1056x1568",
"4:3": "1472x1088",
"3:4": "1088x1472",
"16:9": "1728x960",
"9:16": "960x1728",
},
image_aspect_ratio_param="size",
)

# ── Legacy models (not in current SDK — retained for older SDK versions) ──
Expand Down
Loading
Loading