feat: support multi-image image generation
This commit is contained in:
+10
-1
@@ -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",
|
||||
|
||||
+95
-9
@@ -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",
|
||||
|
||||
+116
-38
@@ -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)
|
||||
)
|
||||
|
||||
|
||||
@@ -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')],
|
||||
},
|
||||
),
|
||||
]
|
||||
@@ -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}"
|
||||
|
||||
+27
-1
@@ -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(
|
||||
|
||||
+146
-1
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user