Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix openai image reference edits

Browse Source
This commit is contained in:
sbyinin
2026-06-19 17:16:09 +08:00
committed by Xubin Ren
parent dd4d410ce2
commit adb737b614
2 changed files with 215 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/providers/image_generation.py
+104 -25
View File
@@ -955,6 +955,56 @@ class OpenAIImageGenerationClient(ImageGenerationProvider):
return model.split("/", 1)[1]
return model
async def _parse_images_response(self, payload: dict[str, Any]) -> list[str]:
client = self._client
owns_client = client is None
if owns_client:
client = httpx.AsyncClient(timeout=self.timeout)
try:
return await _openai_images_from_payload(client, payload)
finally:
if owns_client:
await client.aclose()
async def _post_image_edit(
self,
*,
headers: dict[str, str],
body: dict[str, Any],
reference_images: list[str],
) -> httpx.Response:
files: list[tuple[str, tuple[str, Any, str]]] = []
handles: list[Any] = []
try:
for path in reference_images:
p = Path(path)
raw = p.read_bytes()
mime = detect_image_mime(raw)
if mime is None:
raise ImageGenerationError(f"unsupported reference image: {p}")
handle = p.open("rb")
handles.append(handle)
files.append(("image[]", (p.name, handle, mime)))
client = self._client
if client is not None:
return await client.post(
f"{self.api_base}/images/edits",
headers=headers,
data=body,
files=files,
)
async with httpx.AsyncClient(timeout=self.timeout) as c:
return await c.post(
f"{self.api_base}/images/edits",
headers=headers,
data=body,
files=files,
)
finally:
for handle in handles:
handle.close()
async def generate(
self,
*,
@@ -967,21 +1017,18 @@ class OpenAIImageGenerationClient(ImageGenerationProvider):
if not self.api_key:
raise ImageGenerationError(self.missing_key_message)
if reference_images:
logger.warning(
"DALL-E models do not support reference images; "
"ignoring {} reference image(s) for {}",
len(reference_images),
model,
)
clean_model = self._strip_model_prefix(model)
headers = {
generation_headers = {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
**self.extra_headers,
}
edit_headers = {
"Authorization": f"Bearer {self.api_key}",
**self.extra_headers,
}
clean_model = self._strip_model_prefix(model)
body: dict[str, Any] = {
"model": clean_model,
"prompt": prompt,
@@ -999,13 +1046,37 @@ class OpenAIImageGenerationClient(ImageGenerationProvider):
# Drop null-valued params so extraBody can opt out of defaults like response_format.
body = {key: value for key, value in body.items() if value is not None}
logger.info("OpenAI Images API request: POST {}/images/generations body={}", self.api_base, body)
refs = list(reference_images or [])
if refs:
if not _openai_is_gpt_image_model(clean_model):
raise ImageGenerationError(
f"OpenAI model '{clean_model}' does not support reference images; "
"use a GPT Image model"
)
edit_body = _openai_multipart_form_body(body)
logger.info(
"OpenAI Images API request: POST {}/images/edits body={} reference_images={}",
self.api_base,
edit_body,
len(refs),
)
response = await self._post_image_edit(
headers=edit_headers,
body=edit_body,
reference_images=refs,
)
else:
logger.info(
"OpenAI Images API request: POST {}/images/generations body={}",
self.api_base,
body,
)
response = await self._http_post(
f"{self.api_base}/images/generations",
headers=headers,
body=body,
)
response = await self._http_post(
f"{self.api_base}/images/generations",
headers=generation_headers,
body=body,
)
try:
response.raise_for_status()
@@ -1020,16 +1091,7 @@ class OpenAIImageGenerationClient(ImageGenerationProvider):
logger.info("OpenAI Images API response ({}): {}", response.status_code,
{k: v for k, v in payload.items() if k != "data"})
client = self._client
owns_client = client is None
if owns_client:
client = httpx.AsyncClient(timeout=self.timeout)
try:
images = await _openai_images_from_payload(client, payload)
finally:
if owns_client:
await client.aclose()
images = await self._parse_images_response(payload)
self._require_images(images, payload)
return GeneratedImageResponse(images=images, content="", raw=payload)
@@ -1260,6 +1322,23 @@ def _openai_size(
return "1024x1024"
def _openai_multipart_form_body(body: dict[str, Any]) -> dict[str, str]:
form: dict[str, str] = {}
for key, value in body.items():
if value is None:
continue
if isinstance(value, bool):
form[key] = "true" if value else "false"
elif isinstance(value, str | int | float):
form[key] = str(value)
else:
logger.warning(
"OpenAI image edit parameter '{}' is not a scalar form field; ignoring it",
key,
)
return form
def _openai_is_gpt_image_model(model: str) -> bool:
normalized = model.lower()
return normalized.startswith(("gpt-image", "chatgpt-image"))
tests/providers/test_image_generation.py
+111
View File
@@ -786,6 +786,117 @@ async def test_openai_gpt_image_uses_supported_orientation_sizes() -> None:
assert fake.calls[1]["json"]["size"] == "1536x1024"
@pytest.mark.asyncio
async def test_openai_reference_images_use_edits_endpoint(tmp_path: Path) -> None:
ref = tmp_path / "ref.png"
ref.write_bytes(PNG_BYTES)
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
api_base="https://api.openai.com/v1",
extra_headers={"X-Test": "1"},
client=fake, # type: ignore[arg-type]
)
response = await client.generate(
prompt="make a warmer version",
model="gpt-image-1",
reference_images=[str(ref)],
aspect_ratio="16:9",
)
assert response.images == [PNG_DATA_URL]
call = fake.calls[0]
assert call["url"] == "https://api.openai.com/v1/images/edits"
assert call["headers"]["Authorization"] == "Bearer sk-openai-test"
assert call["headers"]["X-Test"] == "1"
assert "Content-Type" not in call["headers"]
assert "json" not in call
assert call["data"]["model"] == "gpt-image-1"
assert call["data"]["prompt"] == "make a warmer version"
assert call["data"]["size"] == "1536x1024"
assert len(call["files"]) == 1
assert call["files"][0][0] == "image[]"
assert call["files"][0][1][0] == "ref.png"
assert call["files"][0][1][2] == "image/png"
assert call["files"][0][1][1].closed is True
@pytest.mark.asyncio
async def test_openai_reference_images_send_multiple_multipart_files(
tmp_path: Path,
) -> None:
first = tmp_path / "first.png"
second = tmp_path / "second.png"
first.write_bytes(PNG_BYTES)
second.write_bytes(PNG_BYTES)
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
extra_body={
"quality": "high",
"seed": 0,
"safety_checker": False,
"metadata": {"ignored": True},
"background": None,
},
client=fake, # type: ignore[arg-type]
)
await client.generate(
prompt="combine these references",
model="openai/gpt-image-1",
reference_images=[str(first), str(second)],
)
call = fake.calls[0]
assert call["url"] == "https://api.openai.com/v1/images/edits"
assert call["data"]["model"] == "gpt-image-1"
assert call["data"]["prompt"] == "combine these references"
assert call["data"]["quality"] == "high"
assert call["data"]["seed"] == "0"
assert call["data"]["safety_checker"] == "false"
assert "metadata" not in call["data"]
assert "background" not in call["data"]
assert [item[0] for item in call["files"]] == ["image[]", "image[]"]
assert [item[1][0] for item in call["files"]] == ["first.png", "second.png"]
assert all(item[1][1].closed for item in call["files"])
@pytest.mark.asyncio
async def test_openai_gpt_image_without_reference_images_uses_generations_json() -> None:
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
client=fake, # type: ignore[arg-type]
)
await client.generate(prompt="draw", model="gpt-image-1", aspect_ratio="16:9")
call = fake.calls[0]
assert call["url"] == "https://api.openai.com/v1/images/generations"
assert call["headers"]["Content-Type"] == "application/json"
assert call["json"]["model"] == "gpt-image-1"
assert call["json"]["prompt"] == "draw"
assert call["json"]["size"] == "1536x1024"
assert "data" not in call
assert "files" not in call
@pytest.mark.asyncio
async def test_openai_dalle_reference_images_raise_clear_error(tmp_path: Path) -> None:
ref = tmp_path / "ref.png"
ref.write_bytes(PNG_BYTES)
client = OpenAIImageGenerationClient(api_key="sk-openai-test")
with pytest.raises(ImageGenerationError, match="does not support reference images"):
await client.generate(
prompt="edit this",
model="dall-e-3",
reference_images=[str(ref)],
)
@pytest.mark.asyncio
async def test_openai_default_size_when_no_aspect_ratio() -> None:
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))