"""cmhub hosted generation orchestration for AI image studio jobs.""" from __future__ import annotations import os import threading import time import urllib.parse 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) class ImageStudioGenerationError(RuntimeError): """Raised when AI studio hosted generation cannot complete.""" 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 _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 create_generation_jobs( project_id, source_asset_id, prompt, count, *, job_type="main", path=None, ): total = max(0, int(count or 0)) if total <= 0: raise ImageStudioGenerationError("生成数量必须大于0") 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="cmhub", provider="cmhub", 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, ): jobs = create_generation_jobs( project_id, source_asset_id, prompt, count, job_type=job_type, 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, ) 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, ): 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, ) def run_jobs( jobs, *, aspect_ratio="1:1", config=None, cmhub_config_path=None, path=None, should_stop=None, on_event=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": []} 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)) should_stop = should_stop or (lambda: False) lock = threading.Lock() summary = {"total": len(job_list), "success": 0, "failed": 0, "cancelled": 0, "jobs": []} 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, runtime, cfg, image_root, aspect_ratio, path, should_stop, on_event, ): 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": "用户停止"}) 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, runtime, config, image_root, aspect_ratio, db_path, should_stop, on_event): 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": "项目或源图不存在"} try: image_studio.update_job_status(job.id, "running", path=db_path) request_result = _submit_or_resume_job( job, source_asset, runtime, config, aspect_ratio, db_path, should_stop, on_event, ) 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: _raise_if_stopped(should_stop) except ImageStudioGenerationError: try: if os.path.isfile(saved_path): os.remove(saved_path) finally: raise asset = image_studio.add_asset( project.id, _generated_kind(job.job_type), remote_url=request_result["image_url"], local_path=saved_path, aspect_ratio=aspect_ratio, 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 = isinstance(exc, CancelledError) or "停止" in str(exc) status = "cancelled" if cancelled else "failed" error = "用户已停止，已提交任务可稍后继续查询" if cancelled else 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 _submit_or_resume_job( job, source_asset, runtime, config, aspect_ratio, db_path, should_stop, on_event, ): if job.task_id: _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) _raise_if_stopped(should_stop) source_path = _source_path(source_asset) ai_cfg = appconfig.ai_config(config) resolution = str(ai_cfg.get("resolution", "1k") or "1k") payload = { "prompt": str(job.prompt or ""), "model": runtime["alias"], "image_base64": ai._image_data_url(source_path), "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