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

feat(product-suite): add direct job state machine

Browse Source
This commit is contained in:
chengma
2026-07-20 18:12:44 +08:00
parent 2a2cbad7bd
commit aa000ff5e6
11 changed files with 820 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/ai.py
+13 -1
View File
@@ -2776,7 +2776,19 @@ def _save_jpeg(image_bytes, out_path, resolution, jpg_quality):
def _resolution_size(resolution):
return _RESOLUTION_SIZES.get(str(resolution))
text = str(resolution or "").strip().lower()
mapped = _RESOLUTION_SIZES.get(text)
if mapped:
return mapped
if "x" in text:
width, height = text.split("x", 1)
try:
width, height = int(width), int(height)
except (TypeError, ValueError):
return None
if width > 0 and height > 0:
return width, height
return None
def _resolution_size_text(resolution):
app/db.py
+34
View File
@@ -335,6 +335,9 @@ CREATE TABLE IF NOT EXISTS image_studio_assets (
remote_url TEXT,
local_path TEXT,
aspect_ratio TEXT,
requested_output_size TEXT,
rendered_width INTEGER,
rendered_height INTEGER,
parent_asset_id INTEGER REFERENCES image_studio_assets(id) ON DELETE SET NULL,
prompt TEXT,
status TEXT NOT NULL DEFAULT 'available',
@@ -369,6 +372,7 @@ CREATE TABLE IF NOT EXISTS image_studio_jobs (
call_id TEXT,
generation_round_key TEXT,
generation_slot_index INTEGER,
run_session_id TEXT,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL,
submitted_at TEXT,
@@ -508,6 +512,8 @@ def init_db(path=None, conn=None) -> None:
_ensure_image_studio_job_recovery_columns(database)
_ensure_image_studio_generation_round_columns(database)
_ensure_image_studio_job_reference_asset_ids_column(database)
_ensure_image_studio_asset_output_columns(database)
_ensure_image_studio_job_run_session_column(database)
def _ensure_batch_delete_columns(database):
@@ -724,6 +730,34 @@ def _ensure_image_studio_job_reference_asset_ids_column(database):
"ALTER TABLE image_studio_jobs ADD COLUMN reference_asset_ids TEXT"
)
def _ensure_image_studio_asset_output_columns(database):
columns = {
row["name"] for row in database.execute("PRAGMA table_info(image_studio_assets)").fetchall()
}
if "requested_output_size" not in columns:
database.execute(
"ALTER TABLE image_studio_assets ADD COLUMN requested_output_size TEXT"
)
if "rendered_width" not in columns:
database.execute(
"ALTER TABLE image_studio_assets ADD COLUMN rendered_width INTEGER"
)
if "rendered_height" not in columns:
database.execute(
"ALTER TABLE image_studio_assets ADD COLUMN rendered_height INTEGER"
)
def _ensure_image_studio_job_run_session_column(database):
columns = {
row["name"] for row in database.execute("PRAGMA table_info(image_studio_jobs)").fetchall()
}
if "run_session_id" not in columns:
database.execute(
"ALTER TABLE image_studio_jobs ADD COLUMN run_session_id TEXT"
)
def create_batch(file_paths: Iterable[str], note=None, path=None, conn=None) -> str:
batch_id = datetime.now().strftime("%Y%m%d_%H%M%S_") + uuid.uuid4().hex[:8]
files = [os.path.abspath(file_path) for file_path in file_paths]
app/gui/__init__.py
+30 -1
View File
@@ -6,7 +6,7 @@ import os
import sys
from dataclasses import replace
from .. import appconfig, chrome, diagnostics, update_check, update_health
from .. import appconfig, chrome, db, diagnostics, image_studio_generation, update_check, update_health
from ..version import APP_NAME, APP_VERSION, display_name
from . import widgets as _widgets
from .widgets import *
@@ -155,10 +155,39 @@ def main() -> int:
update_health.write_health(health_context, "environment_blocked", str(exc))
QMessageBox.critical(None, "启动配置错误", str(exc))
return 1
runtime_lease = None
try:
database_path = appconfig.db_path(startup["config"])
db.init_db(database_path)
runtime_lease = image_studio_generation.acquire_startup_recovery_lease(
appconfig.data_dir(startup["config"])
)
recovery = image_studio_generation.recover_stale_direct_jobs_at_startup(
lease=runtime_lease,
path=database_path,
)
recovered_count = len(recovery.get("recovered") or [])
if recovered_count:
diagnostics.write_diagnostic_log(
"启动时标记未确认的自定义网关套图任务",
level="WARNING",
step="image_studio_direct_recovery",
payload={"recovered_count": recovered_count},
)
elif recovery.get("skipped"):
diagnostics.write_diagnostic_log(
"检测到另一个程序实例，跳过自定义网关套图任务恢复",
level="INFO",
step="image_studio_direct_recovery",
)
except Exception as exc:
_write_update_check_diagnostic("自定义网关套图任务恢复检查失败", exc=exc)
window = MainWindow(
config=startup["config"],
startup_status=startup["message"],
)
if runtime_lease is not None:
app.aboutToQuit.connect(runtime_lease.release)
window.show()
QTimer.singleShot(
0,
app/gui/workers.py
+11 -9
View File
@@ -262,7 +262,7 @@ class ImageStudioDownloadOriginalWorker(BaseWorker):
class ImageStudioGenerateJobsWorker(BaseWorker):
"""Run cmhub hosted image generation jobs for the AI studio."""
"""Run frozen-source image generation jobs for the AI studio."""
def __init__(
self,
@@ -285,12 +285,13 @@ class ImageStudioGenerateJobsWorker(BaseWorker):
self.job_type = str(job_type or "main")
self.aspect_ratio = str(aspect_ratio or "1:1")
self.db_path = db_path
backend = appconfig.ai_backend(config)
self.config = ai.freeze_runtime_config(
config,
cmhub_config_path=cmhub_config_path,
models_path=(config or {}).get("ai_models_path", appconfig.AI_MODELS_PATH),
include_cmhub=True,
include_direct_models=False,
include_cmhub=backend == "cmhub",
include_direct_models=backend == "direct",
)
self.cmhub_config_path = cmhub_config_path
self._done = 0
@@ -364,12 +365,13 @@ class ProductSuiteGenerateWorker(BaseWorker):
self.generation_round_key = str(generation_round_key or "").strip()
self.aspect_ratio = str(aspect_ratio or "1:1")
self.db_path = db_path
backend = appconfig.ai_backend(config)
self.config = ai.freeze_runtime_config(
config,
cmhub_config_path=cmhub_config_path,
models_path=(config or {}).get("ai_models_path", appconfig.AI_MODELS_PATH),
include_cmhub=True,
include_direct_models=False,
include_cmhub=backend == "cmhub",
include_direct_models=backend == "direct",
)
self.cmhub_config_path = cmhub_config_path
self.job_ids = []
@@ -378,12 +380,11 @@ class ProductSuiteGenerateWorker(BaseWorker):
self._lock = threading.Lock()
def execute(self):
if appconfig.ai_backend(self.config) != "cmhub":
raise ValueError("商品套图仅支持默认网关，请到⑤设置切换后再生成")
total = len(self.job_specs)
if total <= 0:
raise ValueError("商品套图生成任务不能为空")
jobs = []
source = image_studio_generation.generation_source_for_config(self.config)
for spec in self.job_specs:
jobs.append(
image_studio.create_job(
@@ -392,8 +393,8 @@ class ProductSuiteGenerateWorker(BaseWorker):
reference_asset_ids=spec.get("reference_asset_ids"),
job_type=spec.get("job_type") or "套图",
prompt=spec.get("prompt") or "",
generation_source="cmhub",
provider="cmhub",
generation_source=source["generation_source"],
provider=source["provider"],
generation_round_key=spec.get("generation_round_key") or self.generation_round_key or None,
generation_slot_index=spec.get("generation_slot_index"),
path=self.db_path,
@@ -435,6 +436,7 @@ class ProductSuiteGenerateWorker(BaseWorker):
path=self.db_path,
should_stop=self.should_cancel,
on_event=on_event,
run_session_id=self.run_token,
)
summary["project_id"] = self.project_id
summary["job_ids"] = list(self.job_ids)
app/image_studio.py
+115 -6
View File
@@ -33,6 +33,14 @@ JOB_RECOVERY_ACTIONS = {
JOB_RECOVERY_RESUME,
JOB_RECOVERY_REGENERATE,
}
GENERATION_SOURCE_CMHUB = "cmhub"
GENERATION_SOURCE_DIRECT = "direct"
PROVIDER_CMHUB = "cmhub"
PROVIDER_OPENAI_IMAGES_EDITS = "openai_images_edits"
GENERATION_SOURCE_PROVIDERS = {
GENERATION_SOURCE_CMHUB: PROVIDER_CMHUB,
GENERATION_SOURCE_DIRECT: PROVIDER_OPENAI_IMAGES_EDITS,
}
SELECTION_TYPES = {"main", "detail"}
@@ -65,6 +73,9 @@ class ImageStudioAsset:
remote_url: Optional[str]
local_path: Optional[str]
aspect_ratio: Optional[str]
requested_output_size: Optional[str]
rendered_width: Optional[int]
rendered_height: Optional[int]
parent_asset_id: Optional[int]
prompt: Optional[str]
status: str
@@ -95,6 +106,7 @@ class ImageStudioJob:
call_id: Optional[str]
generation_round_key: Optional[str]
generation_slot_index: Optional[int]
run_session_id: Optional[str]
created_at: str
updated_at: str
submitted_at: Optional[str]
@@ -705,6 +717,9 @@ def add_asset(
remote_url=None,
local_path=None,
aspect_ratio=None,
requested_output_size=None,
rendered_width=None,
rendered_height=None,
parent_asset_id=None,
prompt=None,
status=ASSET_STATUS_AVAILABLE,
@@ -720,8 +735,9 @@ def add_asset(
"""
INSERT INTO image_studio_assets
(project_id, kind, remote_url, local_path, aspect_ratio,
requested_output_size, rendered_width, rendered_height,
parent_asset_id, prompt, status, source_order, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
int(project_id),
@@ -729,6 +745,9 @@ def add_asset(
remote_url,
abs_local_path,
aspect_ratio,
str(requested_output_size).strip() if requested_output_size else None,
int(rendered_width) if rendered_width is not None else None,
int(rendered_height) if rendered_height is not None else None,
parent_asset_id,
prompt,
str(status or ASSET_STATUS_AVAILABLE),
@@ -1161,6 +1180,11 @@ def create_job(
path=None,
conn=None,
):
source = str(generation_source or GENERATION_SOURCE_CMHUB).strip().lower()
expected_provider = GENERATION_SOURCE_PROVIDERS.get(source)
provider = str(provider or expected_provider or "").strip().lower()
if expected_provider is None or provider != expected_provider:
raise db.DbError("AI工场任务来源与服务商组合无效")
now = _now()
task_key = str(task_key or _task_key(project_id))
reference_ids = _parse_reference_asset_ids(reference_asset_ids, source_asset_id)
@@ -1200,8 +1224,8 @@ def create_job(
int(project_id),
source_asset_id,
reference_json,
str(generation_source or "cmhub"),
str(provider or "cmhub"),
source,
provider,
str(job_type),
task_key,
str(prompt or ""),
@@ -1664,6 +1688,14 @@ def set_job_submitted(job_id, task_id, *, call_id=None, points_cost=None, points
now = _now()
with _connection(conn, path) as database:
with database:
job = get_job(job_id, conn=database)
if job is None:
raise db.DbError("AI工场任务不存在")
if (
job.generation_source != GENERATION_SOURCE_CMHUB
or job.provider != PROVIDER_CMHUB
):
raise db.DbError("只有默认网关任务可以记录查询任务编号")
database.execute(
"""
UPDATE image_studio_jobs
@@ -1700,6 +1732,7 @@ def update_job_status(
points_balance=None,
recovery_action=None,
increment_attempts=False,
run_session_id=None,
path=None,
conn=None,
):
@@ -1726,6 +1759,7 @@ def update_job_status(
recovery_action = COALESCE(?, recovery_action),
output_asset_id = COALESCE(?, output_asset_id),
points_balance = COALESCE(?, points_balance),
run_session_id = COALESCE(?, run_session_id),
attempts = attempts + ?,
finished_at = CASE WHEN ? THEN ? ELSE finished_at END,
updated_at = ?
@@ -1737,6 +1771,7 @@ def update_job_status(
recovery_action,
output_asset_id,
points_balance,
str(run_session_id).strip() if run_session_id else None,
1 if increment_attempts else 0,
1 if terminal else 0,
now,
@@ -1751,13 +1786,35 @@ def list_resumable_jobs(path=None, conn=None, project_id=None, include_failed_do
if include_failed_downloads:
clauses = [
"task_id IS NOT NULL",
"generation_source = ?",
"provider = ?",
"recovery_action = ?",
"status IN (?, ?, ?, ?)",
]
params = [JOB_RECOVERY_RESUME, "submitted", "running", "failed", "cancelled"]
params = [
GENERATION_SOURCE_CMHUB,
PROVIDER_CMHUB,
JOB_RECOVERY_RESUME,
"submitted",
"running",
"failed",
"cancelled",
]
else:
clauses = ["status IN (?, ?)", "task_id IS NOT NULL", "recovery_action = ?"]
params = ["submitted", "running", JOB_RECOVERY_RESUME]
clauses = [
"status IN (?, ?)",
"task_id IS NOT NULL",
"generation_source = ?",
"provider = ?",
"recovery_action = ?",
]
params = [
"submitted",
"running",
GENERATION_SOURCE_CMHUB,
PROVIDER_CMHUB,
JOB_RECOVERY_RESUME,
]
if project_id is not None:
clauses.append("project_id = ?")
params.append(int(project_id))
@@ -1767,6 +1824,58 @@ def list_resumable_jobs(path=None, conn=None, project_id=None, include_failed_do
return _fetch_all(database, sql, params, ImageStudioJob)
def fail_stale_direct_jobs(*, active_run_session_ids=(), path=None, conn=None):
"""Fail only confirmed stale synchronous direct jobs during application startup."""
active_ids = {
str(value).strip()
for value in (active_run_session_ids or ())
if str(value).strip()
}
with _connection(conn, path) as database:
rows = database.execute(
"""
SELECT * FROM image_studio_jobs
WHERE generation_source = ?
AND provider = ?
AND status = 'running'
ORDER BY id
""",
(GENERATION_SOURCE_DIRECT, PROVIDER_OPENAI_IMAGES_EDITS),
).fetchall()
stale_ids = [
int(row["id"])
for row in rows
if str(row["run_session_id"] or "").strip() not in active_ids
]
if not stale_ids:
return []
now = _now()
with database:
database.executemany(
"""
UPDATE image_studio_jobs
SET status = 'failed',
error = ?,
recovery_action = ?,
finished_at = ?,
updated_at = ?
WHERE id = ?
""",
[
(
"程序中断，无法确认生成结果，请手动重新生成",
JOB_RECOVERY_REGENERATE,
now,
now,
job_id,
)
for job_id in stale_ids
],
)
return [get_job(job_id, conn=database) for job_id in stale_ids]
def replace_selections(project_id, selection_type, asset_ids: Iterable[int], path=None, conn=None):
selection = str(selection_type)
if selection not in SELECTION_TYPES:
app/image_studio_generation.py
+351 -54
View File
@@ -1,4 +1,4 @@
"""cmhub hosted generation orchestration for AI image studio jobs."""
"""Multi-source generation orchestration for AI image studio jobs."""
from __future__ import annotations
@@ -6,21 +6,97 @@ 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_CMHUB_IMAGE_STUDIO_WORKERS = 5
_GLOBAL_IMAGE_STUDIO_SLOTS = threading.BoundedSemaphore(MAX_CMHUB_IMAGE_STUDIO_WORKERS)
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 hosted generation cannot complete."""
"""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):
@@ -36,16 +112,73 @@ 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() == "cmhub"
and str(getattr(job, "provider", "") or "").strip().lower() == "cmhub"
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) != "cmhub":
raise ImageStudioGenerationError("商品套图仅支持默认网关，请到⑤设置切换后再生成")
if appconfig.ai_backend(config) not in {"cmhub", "direct"}:
raise ImageStudioGenerationError("商品套图生成网关配置无效，请到⑤设置检查")
def _generated_kind(job_type):
@@ -104,6 +237,7 @@ def create_generation_jobs(
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(
@@ -112,8 +246,8 @@ def create_generation_jobs(
source_asset_id=source_asset_id,
job_type=job_type,
prompt=prompt,
generation_source="cmhub",
provider="cmhub",
generation_source=source["generation_source"],
provider=source["provider"],
path=path,
)
)
@@ -133,6 +267,7 @@ def generate_image_jobs(
path=None,
should_stop=None,
on_event=None,
run_session_id=None,
):
jobs = create_generation_jobs(
project_id,
@@ -151,6 +286,7 @@ def generate_image_jobs(
path=path,
should_stop=should_stop,
on_event=on_event,
run_session_id=run_session_id,
)
@@ -163,6 +299,7 @@ def resume_image_jobs(
path=None,
should_stop=None,
on_event=None,
run_session_id=None,
):
jobs = image_studio.list_resumable_jobs(
path=path,
@@ -177,6 +314,7 @@ def resume_image_jobs(
path=path,
should_stop=should_stop,
on_event=on_event,
run_session_id=run_session_id,
)
@@ -189,16 +327,18 @@ def run_jobs(
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 = [
job
for job in job_list
if getattr(job, "task_id", None) and not _is_default_gateway_job(job)
]
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),
@@ -209,7 +349,7 @@ def run_jobs(
{
"job": job,
"status": "failed",
"error": "该已提交任务不属于默认网关，不能继续查询",
"error": "AI工场任务来源无效，不能执行",
}
for job in rejected
],
@@ -217,10 +357,23 @@ def run_jobs(
if not job_list:
return summary
runtime = _runtime(cfg, cmhub_config_path)
ai_cfg = appconfig.ai_config(cfg)
image_root = appconfig.image_dir(cfg)
max_workers = min(MAX_CMHUB_IMAGE_STUDIO_WORKERS, max(1, int(ai_cfg.get("image_concurrency", 1) or 1)), len(job_list))
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()
@@ -239,13 +392,14 @@ def run_jobs(
executor.submit(
_run_one_job_with_global_slot,
job.id,
runtime,
runtimes,
cfg,
image_root,
aspect_ratio,
output_spec,
path,
should_stop,
on_event,
run_session_id,
): job
for job in job_list
}
@@ -273,6 +427,7 @@ def run_jobs(
)
futures.pop(future, None)
record({"job": job, "status": "cancelled", "error": "用户停止"})
summary["output"] = dict(output_spec)
return summary
@@ -281,7 +436,17 @@ def _run_one_job_with_global_slot(*args):
return _run_one_job(*args)
def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, should_stop, on_event):
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工场生图任务不存在")
@@ -305,48 +470,94 @@ def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, sho
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", path=db_path)
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)
request_result = _submit_or_resume_job(
job,
source_asset,
runtime,
config,
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)
out_path = _output_path(project, job, image_root)
saved_path = _download_and_save_job_image(
request_result,
out_path,
config,
on_event,
job.id,
should_stop=should_stop,
)
try:
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)
except ImageStudioGenerationError:
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:
if os.path.isfile(saved_path):
os.remove(saved_path)
finally:
raise
_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=request_result["image_url"],
remote_url=remote_url,
local_path=saved_path,
aspect_ratio=aspect_ratio,
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,
@@ -361,9 +572,17 @@ def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, sho
_notify(on_event, {"job_id": job.id, "step": "job_done", "result": "success"})
return {"job": updated, "asset": asset, "status": "succeeded"}
except Exception as exc:
cancelled = isinstance(exc, CancelledError) or "停止" in str(exc)
cancelled = (
not is_direct
and (isinstance(exc, CancelledError) or "停止" in str(exc))
)
status = "cancelled" if cancelled else "failed"
error = "用户已停止，已提交任务可稍后继续查询" if cancelled else str(exc)
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,
@@ -398,6 +617,84 @@ def _reference_assets_for_job(job, db_path):
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,