"""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