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/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)
)