902 lines
30 KiB
Python
902 lines
30 KiB
Python
"""Multi-source generation orchestration for AI image studio jobs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import threading
|
|
import time
|
|
import urllib.parse
|
|
import uuid
|
|
from concurrent.futures import CancelledError, FIRST_COMPLETED, ThreadPoolExecutor, wait
|
|
|
|
from . import ai, appconfig, image_studio
|
|
from .version import APP_VERSION
|
|
|
|
|
|
MAX_IMAGE_STUDIO_WORKERS = 5
|
|
_GLOBAL_IMAGE_STUDIO_SLOTS = threading.BoundedSemaphore(MAX_IMAGE_STUDIO_WORKERS)
|
|
CMHUB_IMAGE_STUDIO_MAX_INPUT_IMAGES = 8
|
|
CMHUB_IMAGE_STUDIO_MAX_SINGLE_INPUT_BYTES = 10 * 1024 * 1024
|
|
CMHUB_IMAGE_STUDIO_MAX_TOTAL_INPUT_BYTES = 32 * 1024 * 1024
|
|
|
|
|
|
class ImageStudioGenerationError(RuntimeError):
|
|
"""Raised when AI studio image generation cannot complete."""
|
|
|
|
|
|
class ImageStudioRuntimeLease:
|
|
"""A process-lifetime lock used only to make startup recovery safe."""
|
|
|
|
def __init__(self, handle, path):
|
|
self._handle = handle
|
|
self.path = path
|
|
self.session_id = "startup-" + uuid.uuid4().hex
|
|
|
|
def release(self):
|
|
handle, self._handle = self._handle, None
|
|
if handle is None:
|
|
return
|
|
try:
|
|
handle.seek(0)
|
|
if os.name == "nt":
|
|
import msvcrt
|
|
|
|
msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1)
|
|
else: # pragma: no cover - Windows release is the supported path.
|
|
import fcntl
|
|
|
|
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
|
|
finally:
|
|
handle.close()
|
|
|
|
|
|
def acquire_startup_recovery_lease(data_dir):
|
|
"""Acquire a non-blocking process lease; return None when another app owns it."""
|
|
|
|
directory = os.path.abspath(str(data_dir or ""))
|
|
if not directory:
|
|
return None
|
|
try:
|
|
os.makedirs(directory, exist_ok=True)
|
|
path = os.path.join(directory, ".image_studio_generation.lock")
|
|
handle = open(path, "a+b")
|
|
if not handle.read(1):
|
|
handle.write(b"0")
|
|
handle.flush()
|
|
handle.seek(0)
|
|
if os.name == "nt":
|
|
import msvcrt
|
|
|
|
msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1)
|
|
else: # pragma: no cover - retained for developer tests outside Windows.
|
|
import fcntl
|
|
|
|
fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
return ImageStudioRuntimeLease(handle, path)
|
|
except (OSError, ImportError):
|
|
try:
|
|
handle.close()
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
def recover_stale_direct_jobs_at_startup(*, lease, path=None):
|
|
"""Recover interrupted direct jobs only while the process lease is held."""
|
|
|
|
if lease is None:
|
|
return {"recovered": [], "skipped": True}
|
|
recovered = image_studio.fail_stale_direct_jobs(path=path)
|
|
return {"recovered": recovered, "skipped": False}
|
|
|
|
|
|
RATIO_OUTPUT_SIZES = {
|
|
"1:1": ("1024x1024", False),
|
|
"3:4": ("1024x1536", True),
|
|
"9:16": ("1024x1536", True),
|
|
"4:3": ("1536x1024", True),
|
|
"16:9": ("1536x1024", True),
|
|
}
|
|
|
|
|
|
def _notify(callback, event):
|
|
if callback is None:
|
|
return
|
|
try:
|
|
callback(dict(event))
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _runtime(config, cmhub_config_path):
|
|
return ai._cmhub_runtime(config, "image", cmhub_config_path)
|
|
|
|
|
|
def _direct_runtime(config):
|
|
ai_cfg = appconfig.ai_config(config)
|
|
models_path = appconfig.ai_models_config_path(config)
|
|
model = ai._role_model(
|
|
"image",
|
|
ai_cfg.get("default_image_model"),
|
|
models_path,
|
|
models=ai._runtime_direct_models(config),
|
|
)
|
|
error = appconfig.image_model_config_error(model)
|
|
if error:
|
|
raise ImageStudioGenerationError(error)
|
|
return {
|
|
"model": model,
|
|
"connect_timeout": ai._connect_timeout(model),
|
|
"read_timeout": ai._read_timeout(model, config),
|
|
"quality": ai._jpg_quality(ai_cfg.get("jpg_quality", 90)),
|
|
}
|
|
|
|
|
|
def _is_default_gateway_job(job):
|
|
return (
|
|
str(getattr(job, "generation_source", "") or "").strip().lower()
|
|
== image_studio.GENERATION_SOURCE_CMHUB
|
|
and str(getattr(job, "provider", "") or "").strip().lower()
|
|
== image_studio.PROVIDER_CMHUB
|
|
)
|
|
|
|
|
|
def _is_direct_gateway_job(job):
|
|
return (
|
|
str(getattr(job, "generation_source", "") or "").strip().lower()
|
|
== image_studio.GENERATION_SOURCE_DIRECT
|
|
and str(getattr(job, "provider", "") or "").strip().lower()
|
|
== image_studio.PROVIDER_OPENAI_IMAGES_EDITS
|
|
)
|
|
|
|
|
|
def generation_source_for_config(config):
|
|
"""Return the persisted source/provider pair for one frozen run config."""
|
|
|
|
if appconfig.ai_backend(config) == "direct":
|
|
return {
|
|
"generation_source": image_studio.GENERATION_SOURCE_DIRECT,
|
|
"provider": image_studio.PROVIDER_OPENAI_IMAGES_EDITS,
|
|
}
|
|
return {
|
|
"generation_source": image_studio.GENERATION_SOURCE_CMHUB,
|
|
"provider": image_studio.PROVIDER_CMHUB,
|
|
}
|
|
|
|
|
|
def requested_output_spec(aspect_ratio):
|
|
"""Map a product-suite ratio to an explicit request size and approximation flag."""
|
|
|
|
ratio = str(aspect_ratio or "1:1")
|
|
size, approximate = RATIO_OUTPUT_SIZES.get(ratio, RATIO_OUTPUT_SIZES["1:1"])
|
|
return {
|
|
"aspect_ratio": ratio,
|
|
"requested_output_size": size,
|
|
"approximate_ratio": bool(approximate),
|
|
}
|
|
|
|
|
|
def _ensure_new_submission_allowed(config):
|
|
if appconfig.ai_backend(config) not in {"cmhub", "direct"}:
|
|
raise ImageStudioGenerationError("商品套图生成网关配置无效，请到⑤设置检查")
|
|
|
|
|
|
def _generated_kind(job_type):
|
|
return "generated_detail" if str(job_type) == "detail" else "generated_main"
|
|
|
|
|
|
def _output_path(project, job, image_root):
|
|
dirs = image_studio.project_image_dirs(image_root, project)
|
|
filename = "job_%06d_%s.jpg" % (int(job.id), _generated_kind(job.job_type))
|
|
return os.path.join(dirs["generated"], filename)
|
|
|
|
|
|
def _source_path(source_asset):
|
|
path = str(getattr(source_asset, "local_path", "") or "").strip()
|
|
if not path or not os.path.isfile(path):
|
|
raise ImageStudioGenerationError("源图尚未下载到本地，不能生成")
|
|
return path
|
|
|
|
|
|
def _build_cmhub_images(source_asset, reference_assets=()):
|
|
"""Build the ordered cmhub image input array from local image assets."""
|
|
candidates = [source_asset] + list(reference_assets or [])
|
|
selected = candidates[:CMHUB_IMAGE_STUDIO_MAX_INPUT_IMAGES]
|
|
omitted_count = max(0, len(candidates) - len(selected))
|
|
images = []
|
|
total_bytes = 0
|
|
for index, asset in enumerate(selected, 1):
|
|
source_path = _source_path(asset)
|
|
file_size = os.path.getsize(source_path)
|
|
if file_size > CMHUB_IMAGE_STUDIO_MAX_SINGLE_INPUT_BYTES:
|
|
role = "主图" if index == 1 else "参考图"
|
|
raise ImageStudioGenerationError("商品套图%s超过10MiB，不能提交" % role)
|
|
data_url = ai._image_data_url(source_path)
|
|
total_bytes += len(data_url.encode("utf-8"))
|
|
if total_bytes > CMHUB_IMAGE_STUDIO_MAX_TOTAL_INPUT_BYTES:
|
|
raise ImageStudioGenerationError(
|
|
"商品套图提交图片总大小超过32MiB，请减少图片或更换较小原图"
|
|
)
|
|
images.append({"image_base64": data_url})
|
|
if not images:
|
|
raise ImageStudioGenerationError("商品套图至少需要一张本地原图")
|
|
return images, omitted_count
|
|
|
|
|
|
def create_generation_jobs(
|
|
project_id,
|
|
source_asset_id,
|
|
prompt,
|
|
count,
|
|
*,
|
|
job_type="main",
|
|
config=None,
|
|
path=None,
|
|
):
|
|
_ensure_new_submission_allowed(config)
|
|
total = max(0, int(count or 0))
|
|
if total <= 0:
|
|
raise ImageStudioGenerationError("生成数量必须大于0")
|
|
source = generation_source_for_config(config)
|
|
jobs = []
|
|
for _ in range(total):
|
|
jobs.append(
|
|
image_studio.create_job(
|
|
project_id,
|
|
source_asset_id=source_asset_id,
|
|
job_type=job_type,
|
|
prompt=prompt,
|
|
generation_source=source["generation_source"],
|
|
provider=source["provider"],
|
|
path=path,
|
|
)
|
|
)
|
|
return jobs
|
|
|
|
|
|
def generate_image_jobs(
|
|
project_id,
|
|
source_asset_id,
|
|
prompt,
|
|
count,
|
|
*,
|
|
job_type="main",
|
|
aspect_ratio="1:1",
|
|
config=None,
|
|
cmhub_config_path=None,
|
|
path=None,
|
|
should_stop=None,
|
|
on_event=None,
|
|
run_session_id=None,
|
|
):
|
|
jobs = create_generation_jobs(
|
|
project_id,
|
|
source_asset_id,
|
|
prompt,
|
|
count,
|
|
job_type=job_type,
|
|
config=config,
|
|
path=path,
|
|
)
|
|
return run_jobs(
|
|
jobs,
|
|
aspect_ratio=aspect_ratio,
|
|
config=config,
|
|
cmhub_config_path=cmhub_config_path,
|
|
path=path,
|
|
should_stop=should_stop,
|
|
on_event=on_event,
|
|
run_session_id=run_session_id,
|
|
)
|
|
|
|
|
|
def resume_image_jobs(
|
|
*,
|
|
project_id=None,
|
|
aspect_ratio="1:1",
|
|
config=None,
|
|
cmhub_config_path=None,
|
|
path=None,
|
|
should_stop=None,
|
|
on_event=None,
|
|
run_session_id=None,
|
|
):
|
|
jobs = image_studio.list_resumable_jobs(
|
|
path=path,
|
|
project_id=project_id,
|
|
include_failed_downloads=True,
|
|
)
|
|
return run_jobs(
|
|
jobs,
|
|
aspect_ratio=aspect_ratio,
|
|
config=config,
|
|
cmhub_config_path=cmhub_config_path,
|
|
path=path,
|
|
should_stop=should_stop,
|
|
on_event=on_event,
|
|
run_session_id=run_session_id,
|
|
)
|
|
|
|
|
|
def run_jobs(
|
|
jobs,
|
|
*,
|
|
aspect_ratio="1:1",
|
|
config=None,
|
|
cmhub_config_path=None,
|
|
path=None,
|
|
should_stop=None,
|
|
on_event=None,
|
|
run_session_id=None,
|
|
):
|
|
cfg = appconfig.load_config() if config is None else config
|
|
job_list = list(jobs or [])
|
|
if not job_list:
|
|
return {"total": 0, "success": 0, "failed": 0, "cancelled": 0, "jobs": []}
|
|
rejected = []
|
|
for job in job_list:
|
|
if _is_direct_gateway_job(job) and getattr(job, "task_id", None):
|
|
rejected.append(job)
|
|
elif not _is_default_gateway_job(job) and not _is_direct_gateway_job(job):
|
|
rejected.append(job)
|
|
job_list = [job for job in job_list if job not in rejected]
|
|
summary = {
|
|
"total": len(job_list) + len(rejected),
|
|
"success": 0,
|
|
"failed": len(rejected),
|
|
"cancelled": 0,
|
|
"jobs": [
|
|
{
|
|
"job": job,
|
|
"status": "failed",
|
|
"error": "AI工场任务来源无效，不能执行",
|
|
}
|
|
for job in rejected
|
|
],
|
|
}
|
|
if not job_list:
|
|
return summary
|
|
|
|
ai_cfg = appconfig.ai_config(cfg)
|
|
image_root = appconfig.image_dir(cfg)
|
|
runtimes = {}
|
|
if any(_is_default_gateway_job(job) for job in job_list):
|
|
runtimes[image_studio.GENERATION_SOURCE_CMHUB] = _runtime(
|
|
cfg,
|
|
cmhub_config_path,
|
|
)
|
|
if any(_is_direct_gateway_job(job) for job in job_list):
|
|
runtimes[image_studio.GENERATION_SOURCE_DIRECT] = _direct_runtime(cfg)
|
|
run_session_id = str(run_session_id or "").strip() or "direct-" + uuid.uuid4().hex
|
|
output_spec = requested_output_spec(aspect_ratio)
|
|
max_workers = min(
|
|
MAX_IMAGE_STUDIO_WORKERS,
|
|
max(1, int(ai_cfg.get("image_concurrency", 1) or 1)),
|
|
len(job_list),
|
|
)
|
|
should_stop = should_stop or (lambda: False)
|
|
lock = threading.Lock()
|
|
|
|
def record(result):
|
|
with lock:
|
|
summary["jobs"].append(result)
|
|
if result.get("status") == "succeeded":
|
|
summary["success"] += 1
|
|
elif result.get("status") == "cancelled":
|
|
summary["cancelled"] += 1
|
|
else:
|
|
summary["failed"] += 1
|
|
|
|
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
|
futures = {
|
|
executor.submit(
|
|
_run_one_job_with_global_slot,
|
|
job.id,
|
|
runtimes,
|
|
cfg,
|
|
image_root,
|
|
output_spec,
|
|
path,
|
|
should_stop,
|
|
on_event,
|
|
run_session_id,
|
|
): job
|
|
for job in job_list
|
|
}
|
|
while futures:
|
|
done, _ = wait(
|
|
set(futures),
|
|
timeout=0.2,
|
|
return_when=FIRST_COMPLETED,
|
|
)
|
|
for future in done:
|
|
futures.pop(future)
|
|
try:
|
|
record(future.result())
|
|
except Exception as exc:
|
|
record({"status": "failed", "error": str(exc)})
|
|
if should_stop():
|
|
for future, job in list(futures.items()):
|
|
if future.cancel():
|
|
image_studio.update_job_status(
|
|
job.id,
|
|
"cancelled",
|
|
error="用户停止",
|
|
recovery_action=_recovery_action_for_job(job),
|
|
path=path,
|
|
)
|
|
futures.pop(future, None)
|
|
record({"job": job, "status": "cancelled", "error": "用户停止"})
|
|
summary["output"] = dict(output_spec)
|
|
return summary
|
|
|
|
|
|
def _run_one_job_with_global_slot(*args):
|
|
with _GLOBAL_IMAGE_STUDIO_SLOTS:
|
|
return _run_one_job(*args)
|
|
|
|
|
|
def _run_one_job(
|
|
job_id,
|
|
runtimes,
|
|
config,
|
|
image_root,
|
|
output_spec,
|
|
db_path,
|
|
should_stop,
|
|
on_event,
|
|
run_session_id,
|
|
):
|
|
job = image_studio.get_job(job_id, path=db_path)
|
|
if job is None:
|
|
raise ImageStudioGenerationError("AI工场生图任务不存在")
|
|
if should_stop():
|
|
updated = image_studio.update_job_status(
|
|
job.id,
|
|
"cancelled",
|
|
error="用户停止",
|
|
recovery_action=_recovery_action_for_job(job),
|
|
path=db_path,
|
|
)
|
|
return {"job": updated, "status": "cancelled", "error": "用户停止"}
|
|
project = image_studio.get_project(job.project_id, path=db_path)
|
|
source_asset = image_studio.get_asset(job.source_asset_id, path=db_path)
|
|
if project is None or source_asset is None:
|
|
updated = image_studio.update_job_status(
|
|
job.id,
|
|
"failed",
|
|
error="项目或源图不存在",
|
|
recovery_action=image_studio.JOB_RECOVERY_REGENERATE,
|
|
path=db_path,
|
|
)
|
|
return {"job": updated, "status": "failed", "error": "项目或源图不存在"}
|
|
is_direct = _is_direct_gateway_job(job)
|
|
if not (_is_default_gateway_job(job) or is_direct):
|
|
raise ImageStudioGenerationError("AI工场任务来源无效，不能执行")
|
|
runtime = runtimes.get(job.generation_source)
|
|
if runtime is None:
|
|
raise ImageStudioGenerationError("当前运行缺少该任务来源的配置快照")
|
|
try:
|
|
image_studio.update_job_status(
|
|
job.id,
|
|
"running",
|
|
run_session_id=run_session_id if is_direct else None,
|
|
path=db_path,
|
|
)
|
|
reference_assets = ()
|
|
if not job.task_id:
|
|
reference_assets = _reference_assets_for_job(job, db_path)
|
|
out_path = _output_path(project, job, image_root)
|
|
if is_direct:
|
|
request_result = _submit_direct_job(
|
|
job,
|
|
source_asset,
|
|
runtime,
|
|
config,
|
|
output_spec,
|
|
should_stop,
|
|
on_event,
|
|
reference_assets=reference_assets,
|
|
)
|
|
# A submitted direct request cannot be cancelled. If bytes arrive after
|
|
# stop was requested, save them before ending the remaining queue.
|
|
saved_path = _save_direct_job_image(
|
|
request_result,
|
|
out_path,
|
|
on_event,
|
|
job.id,
|
|
)
|
|
remote_url = None
|
|
else:
|
|
request_result = _submit_or_resume_job(
|
|
job,
|
|
source_asset,
|
|
runtime,
|
|
config,
|
|
output_spec["aspect_ratio"],
|
|
db_path,
|
|
should_stop,
|
|
on_event,
|
|
reference_assets=reference_assets,
|
|
)
|
|
image_studio.update_job_status(job.id, "running", path=db_path)
|
|
request_result = _poll_job(
|
|
job.id,
|
|
request_result["task_id"],
|
|
runtime,
|
|
request_result,
|
|
db_path,
|
|
should_stop,
|
|
on_event,
|
|
)
|
|
_raise_if_stopped(should_stop)
|
|
request_result["resolution"] = output_spec["requested_output_size"]
|
|
saved_path = _download_and_save_job_image(
|
|
request_result,
|
|
out_path,
|
|
config,
|
|
on_event,
|
|
job.id,
|
|
should_stop=should_stop,
|
|
)
|
|
try:
|
|
_raise_if_stopped(should_stop)
|
|
except ImageStudioGenerationError:
|
|
try:
|
|
if os.path.isfile(saved_path):
|
|
os.remove(saved_path)
|
|
finally:
|
|
raise
|
|
remote_url = request_result["image_url"]
|
|
rendered_width, rendered_height = _rendered_image_size(saved_path)
|
|
asset = image_studio.add_asset(
|
|
project.id,
|
|
_generated_kind(job.job_type),
|
|
remote_url=remote_url,
|
|
local_path=saved_path,
|
|
aspect_ratio=output_spec["aspect_ratio"],
|
|
requested_output_size=output_spec["requested_output_size"],
|
|
rendered_width=rendered_width,
|
|
rendered_height=rendered_height,
|
|
parent_asset_id=source_asset.id,
|
|
prompt=job.prompt,
|
|
path=db_path,
|
|
)
|
|
updated = image_studio.update_job_status(
|
|
job.id,
|
|
"succeeded",
|
|
output_asset_id=asset.id,
|
|
points_balance=request_result.get("points_balance"),
|
|
path=db_path,
|
|
)
|
|
_notify(on_event, {"job_id": job.id, "step": "job_done", "result": "success"})
|
|
return {"job": updated, "asset": asset, "status": "succeeded"}
|
|
except Exception as exc:
|
|
cancelled = (
|
|
not is_direct
|
|
and (isinstance(exc, CancelledError) or "停止" in str(exc))
|
|
)
|
|
status = "cancelled" if cancelled else "failed"
|
|
if cancelled:
|
|
error = "用户已停止，已提交任务可稍后继续查询"
|
|
elif is_direct:
|
|
error = _direct_failure_message(exc)
|
|
else:
|
|
error = str(exc)
|
|
current_job = image_studio.get_job(job.id, path=db_path)
|
|
updated = image_studio.update_job_status(
|
|
job.id,
|
|
status,
|
|
error=error,
|
|
recovery_action=_recovery_action_for_job(current_job),
|
|
path=db_path,
|
|
)
|
|
_notify(
|
|
on_event,
|
|
{
|
|
"job_id": job.id,
|
|
"step": "job_done",
|
|
"result": status,
|
|
"detail": error,
|
|
},
|
|
)
|
|
return {"job": updated, "status": status, "error": error}
|
|
|
|
|
|
def _reference_assets_for_job(job, db_path):
|
|
assets = []
|
|
for asset_id in image_studio.job_reference_asset_ids(job):
|
|
asset = image_studio.get_asset(asset_id, path=db_path)
|
|
if asset is None:
|
|
raise ImageStudioGenerationError("参考图资产不存在，无法继续生成")
|
|
try:
|
|
_source_path(asset)
|
|
except ImageStudioGenerationError as exc:
|
|
raise ImageStudioGenerationError("参考图尚未下载到本地，无法继续生成") from exc
|
|
assets.append(asset)
|
|
return assets
|
|
|
|
|
|
def _direct_image_paths(source_asset, reference_assets=()):
|
|
"""Return ordered local paths; the first source image is the subject image."""
|
|
|
|
return [_source_path(source_asset)] + [
|
|
_source_path(asset) for asset in (reference_assets or ())
|
|
]
|
|
|
|
|
|
def _submit_direct_job(
|
|
job,
|
|
source_asset,
|
|
runtime,
|
|
config,
|
|
output_spec,
|
|
should_stop,
|
|
on_event,
|
|
reference_assets=(),
|
|
):
|
|
if job.task_id:
|
|
raise ImageStudioGenerationError("自定义网关任务不能继续查询，请手动重新生成")
|
|
_raise_if_stopped(should_stop)
|
|
image_paths = _direct_image_paths(source_asset, reference_assets)
|
|
body, content_type = ai._image_edit_body(
|
|
runtime["model"],
|
|
job.prompt,
|
|
image_paths,
|
|
output_spec["requested_output_size"],
|
|
)
|
|
_notify(on_event, {"job_id": job.id, "step": "cover_submit", "result": "start"})
|
|
# Direct images/edits has no idempotency or query endpoint. Retrying an
|
|
# uncertain request could create and charge a second image, so one call only.
|
|
data = ai._call_with_retry(
|
|
runtime["model"],
|
|
body,
|
|
config,
|
|
1,
|
|
request_kind="multipart",
|
|
content_type=content_type,
|
|
)
|
|
image_bytes = ai._extract_image_bytes(data, runtime["model"], config)
|
|
_notify(on_event, {"job_id": job.id, "step": "cover_request", "result": "success"})
|
|
return {
|
|
"image_bytes": image_bytes,
|
|
"resolution": output_spec["requested_output_size"],
|
|
"quality": runtime["quality"],
|
|
}
|
|
|
|
|
|
def _save_direct_job_image(request_result, out_path, on_event, job_id):
|
|
_notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "start"})
|
|
saved_path = ai._save_jpeg(
|
|
request_result["image_bytes"],
|
|
out_path,
|
|
request_result["resolution"],
|
|
request_result["quality"],
|
|
)
|
|
_notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "success"})
|
|
return saved_path
|
|
|
|
|
|
def _rendered_image_size(path):
|
|
try:
|
|
from PIL import Image
|
|
|
|
with Image.open(path) as image:
|
|
width, height = image.size
|
|
return int(width), int(height)
|
|
except Exception as exc:
|
|
raise ImageStudioGenerationError("生成图片保存后无法读取尺寸") from exc
|
|
|
|
|
|
def _direct_failure_message(exc):
|
|
text = str(exc or "").lower()
|
|
if "timeout" in text or "超时" in text:
|
|
return "自定义网关生成超时，无法确认服务商是否已处理，请手动重新生成"
|
|
return "自定义网关生成失败，结果无法确认，请检查图片模型配置后手动重新生成"
|
|
|
|
|
|
def _submit_or_resume_job(
|
|
job,
|
|
source_asset,
|
|
runtime,
|
|
config,
|
|
aspect_ratio,
|
|
db_path,
|
|
should_stop,
|
|
on_event,
|
|
reference_assets=(),
|
|
):
|
|
if job.task_id:
|
|
if not _is_default_gateway_job(job):
|
|
raise ImageStudioGenerationError("该已提交任务不属于默认网关，不能继续查询")
|
|
_notify(on_event, {"job_id": job.id, "step": "cover_request", "result": "resume", "task_id": job.task_id})
|
|
return _request_result(job.task_id, runtime, config)
|
|
_ensure_new_submission_allowed(config)
|
|
_raise_if_stopped(should_stop)
|
|
ai_cfg = appconfig.ai_config(config)
|
|
resolution = str(ai_cfg.get("resolution", "1k") or "1k")
|
|
images, omitted_count = _build_cmhub_images(source_asset, reference_assets)
|
|
if omitted_count:
|
|
_notify(
|
|
on_event,
|
|
{
|
|
"job_id": job.id,
|
|
"step": "cover_submit",
|
|
"result": "warning",
|
|
"detail": "图片最多提交8张，已忽略%d张参考图" % omitted_count,
|
|
},
|
|
)
|
|
payload = {
|
|
"prompt": str(job.prompt or ""),
|
|
"model": runtime["alias"],
|
|
"images": images,
|
|
"resolution": ai._normalize_cmhub_resolution(resolution),
|
|
"aspect_ratio": str(aspect_ratio or "1:1"),
|
|
}
|
|
_notify(on_event, {"job_id": job.id, "step": "cover_submit", "result": "start"})
|
|
data = ai._cmhub_call_with_retry(
|
|
"POST",
|
|
appconfig.cmhub_request_url(runtime["base_url"], "/api/v1/generate/image/tasks"),
|
|
runtime["api_key"],
|
|
payload=payload,
|
|
connect_timeout=runtime["connect_timeout"],
|
|
read_timeout=ai.CMHUB_IMAGE_SUBMIT_READ_TIMEOUT_SECONDS,
|
|
attempts=max(1, int(ai_cfg.get("retry", 0) or 0) + 1),
|
|
headers_extra={
|
|
"Idempotency-Key": job.task_key,
|
|
"X-Client-Version": str(APP_VERSION),
|
|
},
|
|
)
|
|
image_task_id = str(data.get("task_id") or "").strip()
|
|
if not image_task_id:
|
|
raise ImageStudioGenerationError("cmhub 生图任务提交返回格式错误")
|
|
submitted = image_studio.set_job_submitted(
|
|
job.id,
|
|
image_task_id,
|
|
call_id=data.get("call_id"),
|
|
points_cost=data.get("points_cost"),
|
|
points_balance=data.get("points_balance"),
|
|
path=db_path,
|
|
)
|
|
_notify(
|
|
on_event,
|
|
{
|
|
"job_id": job.id,
|
|
"step": "cover_submit",
|
|
"result": "success",
|
|
"task_id": submitted.task_id,
|
|
"call_id": submitted.call_id,
|
|
"points_cost": submitted.points_cost,
|
|
"points_balance": submitted.points_balance,
|
|
},
|
|
)
|
|
return _request_result(image_task_id, runtime, config, points_balance=data.get("points_balance"))
|
|
|
|
|
|
def _request_result(task_id, runtime, config, points_balance=None):
|
|
ai_cfg = appconfig.ai_config(config)
|
|
resolution = str(ai_cfg.get("resolution", "1k") or "1k")
|
|
return {
|
|
"task_id": task_id,
|
|
"connect_timeout": runtime["connect_timeout"],
|
|
"read_timeout": ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS,
|
|
"resolution": resolution,
|
|
"quality": ai._jpg_quality(ai_cfg.get("jpg_quality", 90)),
|
|
"use_system_proxy": runtime["use_system_proxy"],
|
|
"download_with_curl": runtime["download_with_curl"],
|
|
"points_balance": points_balance,
|
|
}
|
|
|
|
|
|
def _poll_job(job_id, task_id, runtime, request_result, db_path, should_stop, on_event):
|
|
poll_url = appconfig.cmhub_request_url(
|
|
runtime["base_url"],
|
|
"/api/v1/generate/image/tasks/%s" % urllib.parse.quote(str(task_id), safe=""),
|
|
)
|
|
started = time.perf_counter()
|
|
deadline = started + max(1, int(ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS))
|
|
poll_index = 0
|
|
while True:
|
|
_raise_if_stopped(should_stop)
|
|
if time.perf_counter() >= deadline:
|
|
raise ImageStudioGenerationError("等待 cmhub 生图任务完成超时，下次可继续查询")
|
|
_notify(on_event, {"job_id": job_id, "step": "cover_poll", "result": "start", "task_id": task_id})
|
|
data = ai._cmhub_call_once(
|
|
"GET",
|
|
poll_url,
|
|
runtime["api_key"],
|
|
payload=None,
|
|
connect_timeout=runtime["connect_timeout"],
|
|
read_timeout=ai.CMHUB_IMAGE_POLL_READ_TIMEOUT_SECONDS,
|
|
headers_extra={"X-Client-Version": str(APP_VERSION)},
|
|
)
|
|
_raise_if_stopped(should_stop)
|
|
status = str(data.get("status") or "").strip().lower()
|
|
if status in {"queued", "running"}:
|
|
_notify(on_event, {"job_id": job_id, "step": "cover_poll", "result": status, "task_id": task_id})
|
|
ai._sleep_cmhub_poll(poll_index, should_stop)
|
|
poll_index += 1
|
|
continue
|
|
if status == "succeeded":
|
|
image_url = ai._extract_cmhub_image_url(data, runtime["base_url"])
|
|
if not image_url:
|
|
raise ImageStudioGenerationError("cmhub 生图任务成功但没有图片地址")
|
|
merged = dict(request_result)
|
|
merged["image_url"] = image_url
|
|
merged["points_balance"] = data.get("points_balance", merged.get("points_balance"))
|
|
_notify(on_event, {"job_id": job_id, "step": "cover_poll", "result": "success", "task_id": task_id})
|
|
return merged
|
|
if status in {"failed", "expired"}:
|
|
error = data.get("error") if isinstance(data.get("error"), dict) else {}
|
|
message = str(error.get("message") or error.get("code") or status)
|
|
image_studio.update_job_status(
|
|
job_id,
|
|
status,
|
|
error=message,
|
|
recovery_action=image_studio.JOB_RECOVERY_REGENERATE,
|
|
path=db_path,
|
|
)
|
|
raise ImageStudioGenerationError(message)
|
|
raise ImageStudioGenerationError("cmhub 生图任务状态返回格式错误")
|
|
|
|
|
|
def _download_and_save_job_image(
|
|
request_result,
|
|
out_path,
|
|
config,
|
|
on_event,
|
|
job_id,
|
|
should_stop=None,
|
|
):
|
|
_raise_if_stopped(should_stop)
|
|
_notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "start"})
|
|
try:
|
|
image_bytes, _ = ai._download_cmhub_image_with_retry(
|
|
request_result["image_url"],
|
|
connect_timeout=request_result["connect_timeout"],
|
|
read_timeout=request_result["read_timeout"],
|
|
use_system_proxy=request_result.get("use_system_proxy", False),
|
|
download_with_curl=request_result.get("download_with_curl", "auto"),
|
|
should_stop=should_stop,
|
|
)
|
|
except CancelledError as exc:
|
|
raise ImageStudioGenerationError(
|
|
"用户已停止，已提交任务可稍后继续查询"
|
|
) from exc
|
|
_raise_if_stopped(should_stop)
|
|
saved_path = ai._save_jpeg(
|
|
image_bytes,
|
|
out_path,
|
|
request_result.get("resolution") or appconfig.ai_config(config).get("resolution", "1k"),
|
|
request_result.get("quality") or appconfig.ai_config(config).get("jpg_quality", 90),
|
|
)
|
|
try:
|
|
_raise_if_stopped(should_stop)
|
|
except ImageStudioGenerationError:
|
|
try:
|
|
if os.path.isfile(saved_path):
|
|
os.remove(saved_path)
|
|
finally:
|
|
raise
|
|
_notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "success"})
|
|
return saved_path
|
|
|
|
|
|
def _raise_if_stopped(should_stop):
|
|
try:
|
|
stopped = bool(should_stop and should_stop())
|
|
except Exception:
|
|
stopped = False
|
|
if stopped:
|
|
raise ImageStudioGenerationError("用户已停止，已提交任务可稍后继续查询")
|
|
|
|
|
|
def _recovery_action_for_job(job):
|
|
if (
|
|
job is not None
|
|
and getattr(job, "task_id", None)
|
|
and getattr(job, "recovery_action", None) == image_studio.JOB_RECOVERY_RESUME
|
|
):
|
|
return image_studio.JOB_RECOVERY_RESUME
|
|
return image_studio.JOB_RECOVERY_REGENERATE
|