feat(image-generation): add StepFun provider support and StepPlan docs
- Add StepFunImageGenerationClient with step-image-edit-2 / step-1x-medium support - Map aspect ratios to StepFun size strings (WxH order) - Add style_reference for step-1x-medium reference-image generation - Register in image gen provider registry (auto-discovered by nanobot.py) - Add 7 unit tests: payload, default size, explicit size, style_reference (1x/non-1x), missing key, no-images - Add StepFun section to docs/image-generation.md with provider config - Add StepPlan (订阅制) subsection with apiBase override example
This commit is contained in:
@@ -14,6 +14,7 @@ from nanobot.providers.image_generation import (
|
||||
ImageGenerationError,
|
||||
MiniMaxImageGenerationClient,
|
||||
OpenRouterImageGenerationClient,
|
||||
StepFunImageGenerationClient,
|
||||
)
|
||||
|
||||
PNG_BYTES = (
|
||||
@@ -387,3 +388,130 @@ async def test_minimax_payload_and_response_with_reference_image(tmp_path: Path)
|
||||
assert body["aspect_ratio"] == "21:9"
|
||||
assert body["subject_reference"][0]["type"] == "character"
|
||||
assert body["subject_reference"][0]["image_file"].startswith("data:image/png;base64,")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# StepFun (阶跃星辰)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stepfun_payload_and_response_with_aspect_ratio() -> None:
|
||||
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
|
||||
client = StepFunImageGenerationClient(
|
||||
api_key="sk-sf-test",
|
||||
api_base="https://api.stepfun.com/v1",
|
||||
extra_headers={"X-Test": "1"},
|
||||
client=fake, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
response = await client.generate(
|
||||
prompt="a cat on the moon",
|
||||
model="step-image-edit-2",
|
||||
aspect_ratio="16:9",
|
||||
)
|
||||
|
||||
assert response.images == [PNG_DATA_URL]
|
||||
call = fake.calls[0]
|
||||
assert call["url"] == "https://api.stepfun.com/v1/images/generations"
|
||||
assert call["headers"]["Authorization"] == "Bearer sk-sf-test"
|
||||
assert call["headers"]["X-Test"] == "1"
|
||||
body = call["json"]
|
||||
assert body["model"] == "step-image-edit-2"
|
||||
assert body["prompt"] == "a cat on the moon"
|
||||
assert body["response_format"] == "b64_json"
|
||||
assert body["n"] == 1
|
||||
assert body["size"] == "1280x800"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stepfun_default_size_when_no_aspect_ratio() -> None:
|
||||
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
|
||||
client = StepFunImageGenerationClient(
|
||||
api_key="sk-sf-test",
|
||||
api_base="https://api.stepfun.com/v1",
|
||||
client=fake, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
await client.generate(prompt="a dog", model="step-image-edit-2")
|
||||
|
||||
body = fake.calls[0]["json"]
|
||||
assert body["size"] == "1024x1024"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stepfun_uses_explicit_image_size() -> None:
|
||||
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
|
||||
client = StepFunImageGenerationClient(
|
||||
api_key="sk-sf-test",
|
||||
api_base="https://api.stepfun.com/v1",
|
||||
client=fake, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
await client.generate(
|
||||
prompt="a bird",
|
||||
model="step-image-edit-2",
|
||||
image_size="1024x1024",
|
||||
)
|
||||
|
||||
body = fake.calls[0]["json"]
|
||||
assert body["size"] == "1024x1024"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stepfun_style_reference_on_1x_model(tmp_path: Path) -> None:
|
||||
"""step-1x-medium supports style_reference for reference-image generation."""
|
||||
ref = tmp_path / "ref.png"
|
||||
ref.write_bytes(PNG_BYTES)
|
||||
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
|
||||
client = StepFunImageGenerationClient(
|
||||
api_key="sk-sf-test",
|
||||
api_base="https://api.stepfun.com/v1",
|
||||
client=fake, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
await client.generate(
|
||||
prompt="in this style",
|
||||
model="step-1x-medium",
|
||||
reference_images=[str(ref)],
|
||||
)
|
||||
|
||||
body = fake.calls[0]["json"]
|
||||
assert "style_reference" in body
|
||||
assert body["style_reference"]["source_url"].startswith("data:image/png;base64,")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stepfun_no_style_reference_on_non_1x_model() -> None:
|
||||
"""step-image-edit-2 does not use style_reference; reference images are ignored."""
|
||||
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
|
||||
client = StepFunImageGenerationClient(
|
||||
api_key="sk-sf-test",
|
||||
api_base="https://api.stepfun.com/v1",
|
||||
client=fake, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
await client.generate(
|
||||
prompt="a flower",
|
||||
model="step-image-edit-2",
|
||||
reference_images=["/tmp/ref.png"],
|
||||
)
|
||||
|
||||
body = fake.calls[0]["json"]
|
||||
assert "style_reference" not in body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stepfun_requires_api_key() -> None:
|
||||
client = StepFunImageGenerationClient(api_key=None)
|
||||
|
||||
with pytest.raises(ImageGenerationError, match="API key"):
|
||||
await client.generate(prompt="draw", model="step-image-edit-2")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stepfun_no_images_raises() -> None:
|
||||
fake = FakeClient(FakeResponse({"data": [{"text": "sorry"}]}))
|
||||
client = StepFunImageGenerationClient(api_key="sk-sf-test", client=fake) # type: ignore[arg-type]
|
||||
|
||||
with pytest.raises(ImageGenerationError, match="returned no images"):
|
||||
await client.generate(prompt="draw", model="step-image-edit-2")
|
||||
|
||||
Reference in New Issue
Block a user