chengma/cmhub
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: support multi-image image generation

Browse Source
This commit is contained in:
QiuSW
2026-07-17 15:55:38 +08:00
parent bd274e26e3
commit 4a0becbb4d
23 changed files with 690 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/ai/providers/base.py
+2
View File
@@ -82,6 +82,7 @@ class ImageGenerationResult:
class MultimodalImage:
data: bytes
mime_type: str = "image/png"
filename: str = "image.png"
class Provider(Protocol):
@@ -108,6 +109,7 @@ class Provider(Protocol):
image: bytes | None = None,
image_mime_type: str = "image/png",
image_filename: str = "image.png",
images: Sequence[MultimodalImage] | None = None,
resolution: str = "1K",
aspect_ratio: str = "1:1",
parameters: Mapping[str, Any] | None = None,
apps/ai/providers/openai_compatible.py
+67 -9
View File
@@ -138,6 +138,7 @@ class ChatCompletionsProvider(BaseHttpProvider):
image: bytes | None = None,
image_mime_type: str = "image/png",
image_filename: str = "image.png",
images: Sequence[MultimodalImage] | None = None,
resolution: str = "1K",
aspect_ratio: str = "1:1",
parameters: Mapping[str, Any] | None = None,
@@ -149,6 +150,7 @@ class ChatCompletionsProvider(BaseHttpProvider):
prompt,
image=image,
image_mime_type=image_mime_type,
images=images,
parameters=parameters,
)
response = self.session.post(
@@ -243,6 +245,7 @@ class GeminiProvider(ChatCompletionsProvider):
image: bytes | None = None,
image_mime_type: str = "image/png",
image_filename: str = "image.png",
images: Sequence[MultimodalImage] | None = None,
resolution: str = "1K",
aspect_ratio: str = "1:1",
parameters: Mapping[str, Any] | None = None,
@@ -254,6 +257,7 @@ class GeminiProvider(ChatCompletionsProvider):
prompt,
image=image,
image_mime_type=image_mime_type,
images=images,
response_modalities=["TEXT", "IMAGE"],
parameters=parameters,
)
@@ -324,6 +328,7 @@ class ImagesGenerationProvider(BaseHttpProvider):
image: bytes | None = None,
image_mime_type: str = "image/png",
image_filename: str = "image.png",
images: Sequence[MultimodalImage] | None = None,
resolution: str = "1K",
aspect_ratio: str = "1:1",
parameters: Mapping[str, Any] | None = None,
@@ -337,8 +342,17 @@ class ImagesGenerationProvider(BaseHttpProvider):
"resolution": resolution,
"n": 1,
}
if image is not None:
payload["image_urls"] = [image_bytes_to_data_url(image, image_mime_type)]
image_inputs = generation_image_inputs(
image=image,
image_mime_type=image_mime_type,
image_filename=image_filename,
images=images,
)
if image_inputs:
payload["image_urls"] = [
image_bytes_to_data_url(item.data, item.mime_type)
for item in image_inputs
]
apply_extra_body(payload, model, parameters)
response = self.session.post(
url,
@@ -376,12 +390,19 @@ class ImagesEditsProvider(BaseHttpProvider):
image: bytes | None = None,
image_mime_type: str = "image/png",
image_filename: str = "image.png",
images: Sequence[MultimodalImage] | None = None,
resolution: str = "1K",
aspect_ratio: str = "1:1",
parameters: Mapping[str, Any] | None = None,
) -> ImageGenerationResult:
validate_model_config(model)
if image is None:
image_inputs = generation_image_inputs(
image=image,
image_mime_type=image_mime_type,
image_filename=image_filename,
images=images,
)
if not image_inputs:
raise AiCapabilityError("images edits provider requires an input image")
url = normalize_api_url(model.url, API_IMAGES_EDITS)
@@ -392,7 +413,15 @@ class ImagesEditsProvider(BaseHttpProvider):
"size": resolution_to_size(resolution),
}
apply_extra_body(data, model, parameters)
files = {"image": (image_filename, image, image_mime_type)}
files: dict[str, tuple[str, bytes, str]] | list[tuple[str, tuple[str, bytes, str]]]
if len(image_inputs) == 1:
item = image_inputs[0]
files = {"image": (item.filename, item.data, item.mime_type)}
else:
files = [
("image", (item.filename, item.data, item.mime_type))
for item in image_inputs
]
response = self.session.post(
url,
headers=self._headers(model),
@@ -443,13 +472,17 @@ def build_chat_image_payload(
*,
image: bytes | None = None,
image_mime_type: str = "image/png",
images: Sequence[MultimodalImage] | None = None,
parameters: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
return build_chat_text_payload(
return build_chat_vision_payload(
model,
prompt,
image=image,
image_mime_type=image_mime_type,
images=generation_image_inputs(
image=image,
image_mime_type=image_mime_type,
images=images,
),
parameters=parameters,
)
@@ -486,12 +519,17 @@ def build_gemini_payload(
*,
image: bytes | None = None,
image_mime_type: str = "image/png",
images: Sequence[MultimodalImage] | None = None,
response_modalities: list[str],
parameters: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
parts: list[dict[str, Any]] = [{"text": prompt}]
if image is not None:
data_url = image_bytes_to_data_url(image, image_mime_type)
for image_input in generation_image_inputs(
image=image,
image_mime_type=image_mime_type,
images=images,
):
data_url = image_bytes_to_data_url(image_input.data, image_input.mime_type)
mime_type, data = split_data_url(data_url)
parts.append({"inlineData": {"mimeType": mime_type, "data": data}})
payload: dict[str, Any] = {
@@ -522,6 +560,26 @@ def build_gemini_vision_payload(
return payload
def generation_image_inputs(
*,
image: bytes | None = None,
image_mime_type: str = "image/png",
image_filename: str = "image.png",
images: Sequence[MultimodalImage] | None = None,
) -> tuple[MultimodalImage, ...]:
if images is not None:
return tuple(images)
if image is None:
return ()
return (
MultimodalImage(
data=image,
mime_type=image_mime_type,
filename=image_filename,
),
)
def apply_extra_body(
payload: dict[str, Any],
model: ResolvedModel,
apps/ai/tests.py
+86
View File
@@ -22,6 +22,7 @@ from apps.ai.providers.openai_compatible import (
ChatCompletionsProvider,
GeminiProvider,
ImagesEditsProvider,
ImagesGenerationProvider,
)
from apps.ai.providers.utils import image_request_timeout, resolution_to_size
@@ -207,6 +208,34 @@ class ChatCompletionsProviderTests(SimpleTestCase):
self.assertEqual(content[0], {"type": "text", "text": "Generate product image"})
self.assertTrue(content[1]["image_url"]["url"].startswith("data:image/jpeg;base64,"))
def test_generate_image_sends_multiple_chat_images_in_order(self):
encoded = base64.b64encode(b"generated-image").decode("ascii")
session = FakeSession(
FakeResponse({"data": [{"b64_json": encoded}]})
)
provider = ChatCompletionsProvider(session=session)
model = ResolvedModel(
name="Nano Banana 2",
url="https://api.vectorengine.ai/v1/chat/completions",
model="gemini-3.1-flash-image-preview",
api_key="test-key",
api_type="chat",
)
provider.generate_image(
"Generate product image",
model,
images=(
MultimodalImage(b"main-image", "image/jpeg", "main.jpg"),
MultimodalImage(b"reference-image", "image/png", "reference.png"),
),
)
content = session.posts[0]["json"]["messages"][0]["content"]
self.assertEqual(content[0]["text"], "Generate product image")
self.assertTrue(content[1]["image_url"]["url"].startswith("data:image/jpeg;base64,"))
self.assertTrue(content[2]["image_url"]["url"].startswith("data:image/png;base64,"))
@override_settings(AI_IMAGE_UPSTREAM_DEADLINE_SECONDS=180)
def test_generate_image_caps_post_and_download_timeouts(self):
session = FakeSession(
@@ -345,6 +374,34 @@ class ChatCompletionsProviderTests(SimpleTestCase):
)
class ImagesGenerationProviderTests(SimpleTestCase):
def test_generate_image_sends_multiple_json_images_in_order(self):
encoded = base64.b64encode(b"generated-image").decode("ascii")
session = FakeSession(FakeResponse({"data": [{"b64_json": encoded}]}))
provider = ImagesGenerationProvider(session=session)
model = ResolvedModel(
name="JSON image provider",
url="https://images.example.test/v1/images/generations",
model="image-model",
api_key="test-key",
api_type="images",
)
provider.generate_image(
"Use the first image as the main product",
model,
images=(
MultimodalImage(b"main-image", "image/jpeg", "main.jpg"),
MultimodalImage(b"reference-image", "image/png", "reference.png"),
),
)
image_urls = session.posts[0]["json"]["image_urls"]
self.assertEqual(len(image_urls), 2)
self.assertTrue(image_urls[0].startswith("data:image/jpeg;base64,"))
self.assertTrue(image_urls[1].startswith("data:image/png;base64,"))
class ImagesEditsProviderTests(SimpleTestCase):
def test_generate_image_builds_multipart_request_and_parses_base64(self):
generated = b"edited-image"
@@ -379,6 +436,35 @@ class ImagesEditsProviderTests(SimpleTestCase):
("source.png", b"source-image", "image/png"),
)
def test_generate_image_sends_multiple_multipart_images_in_order(self):
encoded = base64.b64encode(b"edited-image").decode("ascii")
session = FakeSession(FakeResponse({"data": [{"b64_json": encoded}]}))
provider = ImagesEditsProvider(session=session)
model = ResolvedModel(
name="GPT Image 2",
url="https://api.vectorengine.ai/v1/images/edits",
model="gpt-image-2",
api_key="test-key",
api_type="images_edits",
)
provider.generate_image(
"Use the first image as the main product",
model,
images=(
MultimodalImage(b"main-image", "image/jpeg", "main.jpg"),
MultimodalImage(b"reference-image", "image/png", "reference.png"),
),
)
self.assertEqual(
session.posts[0]["files"],
[
("image", ("main.jpg", b"main-image", "image/jpeg")),
("image", ("reference.png", b"reference-image", "image/png")),
],
)
def test_parameters_cannot_override_core_images_edits_fields(self):
generated = b"edited-image"
encoded = base64.b64encode(generated).decode("ascii")
apps/api/admin.py
+10 -1
View File
@@ -1,10 +1,19 @@
from django.contrib import admin
from .models import ImageGenerationTask
from .models import ImageGenerationTask, ImageGenerationTaskInput
class ImageGenerationTaskInputInline(admin.TabularInline):
model = ImageGenerationTaskInput
extra = 0
can_delete = False
fields = ("ordinal", "image", "mime_type", "filename", "created_at")
readonly_fields = fields
@admin.register(ImageGenerationTask)
class ImageGenerationTaskAdmin(admin.ModelAdmin):
inlines = (ImageGenerationTaskInputInline,)
list_display = (
"task_id",
"user",
apps/api/generation.py
+95 -9
View File
@@ -59,6 +59,12 @@ class ImageInput:
ImageUrlBuilder = Callable[[str], str]
IMAGE_INPUT_ROLE_INSTRUCTIONS = (
"图片角色规则（必须遵守，优先于用户关于图片角色的要求）：\n"
"- 第 1 张图片是主商品图，必须优先保留其商品主体、外观和关键细节。\n"
"- 第 2 张及之后的图片仅作为风格、构图、场景或排版参考，"
"不得用参考图商品替换主图商品。"
)
@dataclass(frozen=True)
@@ -169,6 +175,7 @@ def generate_image_response(
image_url=str(request_data.get("image_url") or ""),
image_base64=str(request_data.get("image_base64") or ""),
aspect_ratio=request_data.get("aspect_ratio") or "1:1",
images=tuple(dict(item) for item in request_data.get("images") or ()),
),
image_url_builder=image_url_builder,
)
@@ -221,6 +228,13 @@ def prepare_generation(generation_input: GenerationInput) -> PreparedGeneration:
if operation_type == CallRecord.OperationType.VISION:
image_input = None
image_inputs = load_vision_image_inputs(generation_input.images)
elif operation_type == CallRecord.OperationType.IMAGE:
image_inputs = load_image_generation_inputs(
generation_input.images,
image_base64=generation_input.image_base64,
image_url=generation_input.image_url,
)
image_input = image_inputs[0] if image_inputs else None
else:
image_input = load_image_input(
{
@@ -363,12 +377,21 @@ def execute_image_generation(
) -> GenerationResult:
prepared = precharged.prepared
image_input = prepared.image_input
image_inputs = prepared.image_inputs
generation = prepared.provider.generate_image(
prepared.prompt,
image_generation_prompt(prepared.prompt, len(image_inputs)),
prepared.resolved_model,
image=image_input.data if image_input else None,
image_mime_type=image_input.mime_type if image_input else "image/png",
image_filename=image_input.filename if image_input else "image.png",
images=tuple(
MultimodalImage(
data=item.data,
mime_type=item.mime_type,
filename=item.filename,
)
for item in image_inputs
),
resolution=prepared.resolution,
aspect_ratio=prepared.aspect_ratio,
parameters=prepared.parameters,
@@ -530,13 +553,21 @@ def upstream_error(exc: Exception) -> ApiRequestError:
def load_image_input(data: Mapping[str, Any]) -> ImageInput | None:
return load_image_input_with_limit(data)
def load_image_input_with_limit(
data: Mapping[str, Any],
*,
max_bytes: int | None = None,
) -> ImageInput | None:
raw_base64 = str(data.get("image_base64") or "").strip()
if raw_base64:
return decode_image_input(raw_base64)
return decode_image_input(raw_base64, max_bytes=max_bytes)
image_url = str(data.get("image_url") or "").strip()
if image_url:
return download_image_input(image_url)
return download_image_input(image_url, max_bytes=max_bytes)
return None
@@ -544,15 +575,54 @@ def load_image_input(data: Mapping[str, Any]) -> ImageInput | None:
def load_vision_image_inputs(
items: Sequence[Mapping[str, Any]],
) -> tuple[ImageInput, ...]:
max_images = max(1, int(getattr(settings, "VISION_MAX_IMAGES", 8)))
return load_ordered_image_inputs(
items,
max_images=max(1, int(getattr(settings, "VISION_MAX_IMAGES", 8))),
max_image_bytes=max(
1,
int(getattr(settings, "VISION_MAX_IMAGE_BYTES", 10 * 1024 * 1024)),
),
max_total_bytes=max(
1,
int(getattr(settings, "VISION_MAX_TOTAL_BYTES", 32 * 1024 * 1024)),
),
)
def load_image_generation_inputs(
items: Sequence[Mapping[str, Any]],
*,
image_base64: str = "",
image_url: str = "",
) -> tuple[ImageInput, ...]:
max_image_bytes = max(
1,
int(getattr(settings, "VISION_MAX_IMAGE_BYTES", 10 * 1024 * 1024)),
int(getattr(settings, "IMAGE_MAX_INPUT_IMAGE_BYTES", 10 * 1024 * 1024)),
)
max_total_bytes = max(
1,
int(getattr(settings, "VISION_MAX_TOTAL_BYTES", 32 * 1024 * 1024)),
if items:
return load_ordered_image_inputs(
items,
max_images=max(1, int(getattr(settings, "IMAGE_MAX_INPUT_IMAGES", 8))),
max_image_bytes=max_image_bytes,
max_total_bytes=max(
1,
int(getattr(settings, "IMAGE_MAX_INPUT_TOTAL_BYTES", 32 * 1024 * 1024)),
),
)
image_input = load_image_input_with_limit(
{"image_base64": image_base64, "image_url": image_url},
max_bytes=max_image_bytes,
)
return (image_input,) if image_input is not None else ()
def load_ordered_image_inputs(
items: Sequence[Mapping[str, Any]],
*,
max_images: int,
max_image_bytes: int,
max_total_bytes: int,
) -> tuple[ImageInput, ...]:
if not items:
raise ApiRequestError("bad_request", "images 至少需要一张图片", status.HTTP_400_BAD_REQUEST)
if len(items) > max_images:
@@ -591,6 +661,12 @@ def load_vision_image_inputs(
return tuple(image_inputs)
def image_generation_prompt(prompt: str, image_count: int) -> str:
if image_count < 1:
return prompt
return f"{prompt}\n\n{IMAGE_INPUT_ROLE_INSTRUCTIONS}"
def decode_image_input(value: str, *, max_bytes: int | None = None) -> ImageInput:
mime_type = "image/png"
encoded = value
@@ -646,7 +722,10 @@ def download_image_input(url: str, *, max_bytes: int | None = None) -> ImageInpu
content_type = response.headers.get("Content-Type", "image/png").split(";", 1)[0].strip().lower()
if not content_type.startswith("image/"):
raise ApiRequestError("bad_request", "image_url 不是图片资源", status.HTTP_400_BAD_REQUEST)
image = read_limited_image_response(response, max_bytes=max_bytes)
image = read_limited_image_response(
response,
max_bytes=effective_image_url_max_bytes(max_bytes),
)
if not image:
raise ApiRequestError("bad_request", "image_url 图片内容为空", status.HTTP_400_BAD_REQUEST)
return ImageInput(
@@ -746,6 +825,13 @@ def read_limited_image_response(response, *, max_bytes: int | None = None) -> by
return b"".join(chunks)
def effective_image_url_max_bytes(max_bytes: int | None) -> int:
url_limit = max(1, int(getattr(settings, "IMAGE_URL_MAX_BYTES", 10 * 1024 * 1024)))
if max_bytes is None:
return url_limit
return min(url_limit, max(1, int(max_bytes)))
def filename_for_mime(mime_type: str) -> str:
extension = {
"image/jpeg": "jpg",
apps/api/image_tasks.py
+116 -38
View File
@@ -25,7 +25,7 @@ from .generation import (
prepare_generation,
precharge_generation,
)
from .models import ImageGenerationTask
from .models import ImageGenerationTask, ImageGenerationTaskInput
IDEMPOTENCY_KEY_MAX_LENGTH = 128
@@ -61,6 +61,7 @@ def create_image_generation_task(
image_url=str(request_data.get("image_url") or ""),
image_base64=str(request_data.get("image_base64") or ""),
aspect_ratio=request_data.get("aspect_ratio") or "1:1",
images=tuple(dict(item) for item in request_data.get("images") or ()),
)
)
@@ -85,12 +86,16 @@ def create_image_generation_task(
request_hash=request_hash,
request_payload=task_request_payload(
request_data=request_data,
image_input=prepared.image_input,
image_inputs=prepared.image_inputs,
),
points_balance_after_charge=precharged.points_balance_after_charge,
expires_at=now + timedelta(hours=image_task_retention_hours()),
)
store_task_input_image(task, prepared.image_input)
store_task_input_images(
task,
image_inputs=prepared.image_inputs,
request_data=request_data,
)
return task, True
except IntegrityError:
if idempotency_key_hash:
@@ -137,64 +142,111 @@ def normalize_idempotency_key(value: str) -> str:
def request_hash_for_image_request(request_data: Mapping[str, Any]) -> str:
image_base64 = str(request_data.get("image_base64") or "")
payload = {
"prompt": str(request_data.get("prompt") or ""),
"model": str(request_data.get("model") or ""),
"resolution": str(request_data.get("resolution") or "1K"),
"aspect_ratio": str(request_data.get("aspect_ratio") or "1:1"),
"parameters": dict(request_data.get("parameters") or {}),
"image_url": str(request_data.get("image_url") or ""),
"image_base64_sha256": hash_text(image_base64) if image_base64 else "",
}
image_items = request_data.get("images") or ()
if image_items:
payload["images"] = image_request_input_hash_data(request_data)
else:
image_base64 = str(request_data.get("image_base64") or "")
payload["image_url"] = str(request_data.get("image_url") or "")
payload["image_base64_sha256"] = hash_text(image_base64) if image_base64 else ""
return hash_json(payload)
def task_request_payload(
*,
request_data: Mapping[str, Any],
image_input: ImageInput | None,
image_inputs: tuple[ImageInput, ...],
) -> dict[str, Any]:
image_source = "none"
if request_data.get("image_base64"):
image_source = "base64_stored"
elif request_data.get("image_url"):
image_source = "url_stored"
payload = {
return {
"prompt": str(request_data.get("prompt") or ""),
"model": str(request_data.get("model") or ""),
"resolution": str(request_data.get("resolution") or "1K"),
"aspect_ratio": str(request_data.get("aspect_ratio") or "1:1"),
"parameters": dict(request_data.get("parameters") or {}),
"image_url": str(request_data.get("image_url") or ""),
"image_input": None,
"image_inputs": [
{
"source": source,
"mime_type": image_input.mime_type,
"filename": image_input.filename,
"storage_path": "",
}
for image_input, source in zip(
image_inputs,
image_request_input_sources(request_data),
strict=True,
)
],
}
if image_input is not None:
payload["image_input"] = {
"source": image_source,
"mime_type": image_input.mime_type,
"filename": image_input.filename,
"storage_path": "",
def image_request_input_hash_data(request_data: Mapping[str, Any]) -> list[dict[str, str]]:
items = request_data.get("images") or ()
if items:
return [
{
"image_url": str(item.get("image_url") or ""),
"image_base64_sha256": hash_text(str(item.get("image_base64") or ""))
if item.get("image_base64")
else "",
}
for item in items
]
image_base64 = str(request_data.get("image_base64") or "")
return [
{
"image_url": str(request_data.get("image_url") or ""),
"image_base64_sha256": hash_text(image_base64) if image_base64 else "",
}
return payload
]
def store_task_input_image(
def image_request_input_sources(request_data: Mapping[str, Any]) -> tuple[str, ...]:
items = request_data.get("images") or ()
if items:
return tuple(
"base64_stored" if item.get("image_base64") else "url_stored"
for item in items
)
if request_data.get("image_base64"):
return ("base64_stored",)
if request_data.get("image_url"):
return ("url_stored",)
return ()
def store_task_input_images(
task: ImageGenerationTask,
image_input: ImageInput | None,
*,
image_inputs: tuple[ImageInput, ...],
request_data: Mapping[str, Any],
) -> None:
if image_input is None:
if not image_inputs:
return
task.input_image.save(
image_input.filename,
ContentFile(image_input.data),
save=True,
)
payload = dict(task.request_payload or {})
image_info = dict(payload.get("image_input") or {})
image_info["storage_path"] = task.input_image.name
payload["image_input"] = image_info
image_metadata = list(payload.get("image_inputs") or [])
for ordinal, image_input in enumerate(image_inputs):
task_input = ImageGenerationTaskInput(
task=task,
ordinal=ordinal,
mime_type=image_input.mime_type,
filename=image_input.filename,
)
task_input.image.save(
image_input.filename,
ContentFile(image_input.data),
save=False,
)
task_input.save()
image_metadata[ordinal]["storage_path"] = task_input.image.name
payload["image_inputs"] = image_metadata
task.request_payload = payload
task.save(update_fields=("request_payload", "updated_at"))
@@ -332,9 +384,13 @@ def precharged_generation_for_task(task: ImageGenerationTask) -> PrechargedGener
aspect_ratio=str(payload.get("aspect_ratio") or "1:1"),
)
)
image_input = stored_image_input_for_task(task)
if image_input is not None:
prepared = replace(prepared, image_input=image_input)
image_inputs = stored_image_inputs_for_task(task)
if image_inputs:
prepared = replace(
prepared,
image_input=image_inputs[0],
image_inputs=image_inputs,
)
return PrechargedGeneration(
prepared=prepared,
@@ -344,7 +400,28 @@ def precharged_generation_for_task(task: ImageGenerationTask) -> PrechargedGener
)
def stored_image_input_for_task(task: ImageGenerationTask) -> ImageInput | None:
def stored_image_inputs_for_task(task: ImageGenerationTask) -> tuple[ImageInput, ...]:
task_inputs = list(task.input_images.all())
if task_inputs:
return tuple(
ImageInput(
data=read_task_input_image(task_input),
mime_type=task_input.mime_type,
filename=task_input.filename,
)
for task_input in task_inputs
)
legacy_input = stored_legacy_image_input_for_task(task)
return (legacy_input,) if legacy_input is not None else ()
def read_task_input_image(task_input: ImageGenerationTaskInput) -> bytes:
with task_input.image.open("rb") as image_file:
return image_file.read()
def stored_legacy_image_input_for_task(task: ImageGenerationTask) -> ImageInput | None:
if not task.input_image:
return None
payload = dict(task.request_payload or {})
@@ -589,6 +666,7 @@ def refund_task_call(task: ImageGenerationTask, error_message: str) -> None:
def refresh_task(task: ImageGenerationTask) -> ImageGenerationTask:
return (
ImageGenerationTask.objects.select_related("user", "api_key", "call_record")
.prefetch_related("input_images")
.get(pk=task.pk)
)
apps/api/migrations/0003_image_generation_task_input.py
+33
View File
@@ -0,0 +1,33 @@
# Generated by Django 5.2.15 on 2026-07-17 07:29
import django.db.models.deletion
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
('api', '0002_imagegenerationtask_next_attempt_at_and_more'),
]
operations = [
migrations.CreateModel(
name='ImageGenerationTaskInput',
fields=[
('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
('ordinal', models.PositiveSmallIntegerField(verbose_name='输入顺序')),
('image', models.FileField(upload_to='generated/task_inputs/%Y/%m/%d/', verbose_name='输入图片')),
('mime_type', models.CharField(max_length=100, verbose_name='MIME 类型')),
('filename', models.CharField(max_length=255, verbose_name='原始文件名')),
('created_at', models.DateTimeField(auto_now_add=True, verbose_name='创建时间')),
('task', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='input_images', to='api.imagegenerationtask', verbose_name='图片生成任务')),
],
options={
'verbose_name': '图片生成任务输入',
'verbose_name_plural': '图片生成任务输入',
'db_table': 'image_generation_task_input',
'ordering': ('ordinal', 'id'),
'constraints': [models.UniqueConstraint(fields=('task', 'ordinal'), name='unique_image_task_input_ordinal')],
},
),
]
apps/api/models.py
+32
View File
@@ -98,3 +98,35 @@ class ImageGenerationTask(models.Model):
def __str__(self) -> str:
return f"{self.task_id} {self.status}"
class ImageGenerationTaskInput(models.Model):
task = models.ForeignKey(
ImageGenerationTask,
verbose_name="图片生成任务",
on_delete=models.CASCADE,
related_name="input_images",
)
ordinal = models.PositiveSmallIntegerField("输入顺序")
image = models.FileField(
"输入图片",
upload_to="generated/task_inputs/%Y/%m/%d/",
)
mime_type = models.CharField("MIME 类型", max_length=100)
filename = models.CharField("原始文件名", max_length=255)
created_at = models.DateTimeField("创建时间", auto_now_add=True)
class Meta:
db_table = "image_generation_task_input"
verbose_name = "图片生成任务输入"
verbose_name_plural = "图片生成任务输入"
ordering = ("ordinal", "id")
constraints = [
models.UniqueConstraint(
fields=("task", "ordinal"),
name="unique_image_task_input_ordinal",
)
]
def __str__(self) -> str:
return f"{self.task.task_id} #{self.ordinal + 1}"
apps/api/serializers.py
+27 -1
View File
@@ -35,6 +35,7 @@ class GenerateImageRequestSerializer(serializers.Serializer):
)
image_url = serializers.URLField(required=False, allow_blank=True)
image_base64 = serializers.CharField(required=False, allow_blank=True)
images = serializers.ListField(required=False, write_only=True)
resolution = serializers.CharField(
required=False,
allow_blank=True,
@@ -49,8 +50,29 @@ class GenerateImageRequestSerializer(serializers.Serializer):
)
parameters = serializers.DictField(required=False, default=dict)
def validate_images(self, value):
if not value:
raise serializers.ValidationError("images 至少需要一张图片。")
max_images = max(1, int(settings.IMAGE_MAX_INPUT_IMAGES))
if len(value) > max_images:
raise serializers.ValidationError(f"单次最多上传 {max_images} 张图片。")
class VisionImageInputSerializer(serializers.Serializer):
serializer = ImageInputSerializer(data=value, many=True)
serializer.is_valid(raise_exception=True)
return serializer.validated_data
def validate(self, attrs):
has_legacy_image = bool(str(attrs.get("image_url") or "").strip()) or bool(
str(attrs.get("image_base64") or "").strip()
)
if attrs.get("images") is not None and has_legacy_image:
raise serializers.ValidationError(
"images 不能与 image_url 或 image_base64 同时提供。"
)
return attrs
class ImageInputSerializer(serializers.Serializer):
image_url = serializers.URLField(required=False, allow_blank=True)
image_base64 = serializers.CharField(required=False, allow_blank=True)
@@ -64,6 +86,10 @@ class VisionImageInputSerializer(serializers.Serializer):
return attrs
class VisionImageInputSerializer(ImageInputSerializer):
pass
class AnalyzeImagesRequestSerializer(serializers.Serializer):
prompt = serializers.CharField(trim_whitespace=True, allow_blank=False)
model = serializers.CharField(
apps/api/tests.py
+146 -1
View File
@@ -36,7 +36,7 @@ from apps.api.image_tasks import (
reap_stale_image_tasks,
run_image_generation_task,
)
from apps.api.models import ImageGenerationTask
from apps.api.models import ImageGenerationTask, ImageGenerationTaskInput
from apps.api.throttles import GenerateRateThrottle
from apps.api.views import ClientLatestReleaseView, ExternalApiView, ModelsView
from apps.ai.models import AiModel, ModelAlias
@@ -1672,6 +1672,112 @@ class GenerateApiTests(TestCase):
self.assertEqual(call.result_summary, "image_bytes=21")
self.assertNotIn("SECRET_RAW", call.result_ref + call.result_summary)
def test_generate_image_accepts_ordered_images_and_injects_role_rules(self):
first = base64.b64encode(b"main-image").decode("ascii")
second = base64.b64encode(b"reference-image").decode("ascii")
response = self.post_with_provider(
"/api/v1/generate/image",
{
"prompt": "生成新的商品主图",
"model": self.image_alias,
"images": [
{"image_base64": f"data:image/jpeg;base64,{first}"},
{"image_base64": f"data:image/png;base64,{second}"},
],
},
)
self.assertEqual(response.status_code, 200)
self.assertEqual(response.data["points_cost"], 10)
self.assertEqual(response.data["points_balance"], 90)
provider_call = self.provider.image_calls[0]
self.assertEqual(provider_call["image"], b"main-image")
self.assertEqual(
[image.data for image in provider_call["images"]],
[b"main-image", b"reference-image"],
)
self.assertIn("第 1 张图片是主商品图", provider_call["prompt"])
self.assertIn("第 2 张及之后的图片仅作为", provider_call["prompt"])
self.assertIn("生成新的商品主图", provider_call["prompt"])
def test_generate_image_accepts_mixed_base64_and_url_images_in_order(self):
encoded = base64.b64encode(b"main-image").decode("ascii")
downloaded = FakeImageUrlResponse(
headers={"Content-Type": "image/jpeg"},
chunks=(b"reference-image",),
)
with (
patch(
"apps.api.generation.socket.getaddrinfo",
return_value=dns_result("93.184.216.34"),
),
patch(
"apps.api.generation.requests.Session.get",
return_value=downloaded,
),
):
response = self.post_with_provider(
"/api/v1/generate/image",
{
"prompt": "生成新的商品主图",
"model": self.image_alias,
"images": [
{"image_base64": f"data:image/png;base64,{encoded}"},
{"image_url": "https://images.example.test/reference.jpg"},
],
},
)
self.assertEqual(response.status_code, 200)
provider_call = self.provider.image_calls[0]
self.assertEqual(
[image.data for image in provider_call["images"]],
[b"main-image", b"reference-image"],
)
self.assertEqual(
[image.mime_type for image in provider_call["images"]],
["image/png", "image/jpeg"],
)
def test_generate_image_rejects_mixed_legacy_and_images_inputs_without_charge(self):
encoded = base64.b64encode(b"main-image").decode("ascii")
response = self.post_with_provider(
"/api/v1/generate/image",
{
"prompt": "生成图片",
"model": self.image_alias,
"image_base64": f"data:image/png;base64,{encoded}",
"images": [{"image_base64": f"data:image/png;base64,{encoded}"}],
},
)
self.assertEqual(response.status_code, 400)
self.assertEqual(self.provider.image_calls, [])
self.assert_generation_not_charged()
@override_settings(IMAGE_MAX_INPUT_IMAGES=1)
def test_generate_image_rejects_too_many_input_images_without_charge(self):
encoded = base64.b64encode(b"input-image").decode("ascii")
response = self.post_with_provider(
"/api/v1/generate/image",
{
"prompt": "生成图片",
"model": self.image_alias,
"images": [
{"image_base64": encoded},
{"image_base64": encoded},
],
},
)
self.assertEqual(response.status_code, 400)
self.assertEqual(self.provider.image_calls, [])
self.assert_generation_not_charged()
def test_sync_image_usage_telemetry_logs_safe_client_version_and_key_identity(self):
encoded = base64.b64encode(b"input-image").decode("ascii")
payload = {
@@ -1883,6 +1989,45 @@ class GenerateApiTests(TestCase):
self.assertEqual(poll.data["result"]["image_url"], task.result_url)
self.assertEqual(repeat.data["result"]["image_url"], task.result_url)
@override_settings(MEDIA_PUBLIC_BASE_URL="https://cm.example.test")
def test_async_image_task_stores_and_restores_ordered_inputs(self):
first = base64.b64encode(b"main-image").decode("ascii")
second = base64.b64encode(b"reference-image").decode("ascii")
response = self.post_with_provider(
"/api/v1/generate/image/tasks",
{
"prompt": "生成新的商品主图",
"model": self.image_alias,
"images": [
{"image_base64": f"data:image/jpeg;base64,{first}"},
{"image_base64": f"data:image/png;base64,{second}"},
],
},
)
self.assertEqual(response.status_code, 202)
task = ImageGenerationTask.objects.get(task_id=response.data["task_id"])
stored_inputs = list(task.input_images.order_by("ordinal"))
self.assertEqual(len(stored_inputs), 2)
self.assertEqual([item.ordinal for item in stored_inputs], [0, 1])
self.assertFalse(bool(task.input_image))
self.assertFalse(ImageGenerationTaskInput.objects.filter(task=task, image__isnull=True).exists())
serialized = json.dumps(task.request_payload, ensure_ascii=False)
self.assertNotIn(first, serialized)
self.assertNotIn(second, serialized)
with patch("apps.api.generation.get_provider", return_value=self.provider):
claimed = claim_next_image_task("worker-multi")
completed = run_image_generation_task(claimed, worker_id="worker-multi")
self.assertEqual(completed.status, ImageGenerationTask.Status.SUCCEEDED)
provider_call = self.provider.image_calls[0]
self.assertEqual(
[image.data for image in provider_call["images"]],
[b"main-image", b"reference-image"],
)
self.assertIn("第 1 张图片是主商品图", provider_call["prompt"])
def test_async_image_poll_rejects_cross_user_access(self):
response = self.post_with_provider(
"/api/v1/generate/image/tasks",