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

feat: snapshot product suite reference assets

Browse Source
This commit is contained in:
chengma
2026-07-17 17:03:40 +08:00
parent d11dc752b5
commit 1fabd8afcf
12 changed files with 252 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/db.py
+12
View File
@@ -338,6 +338,7 @@ CREATE TABLE IF NOT EXISTS image_studio_jobs (
id INTEGER PRIMARY KEY,
project_id INTEGER NOT NULL REFERENCES image_studio_projects(id) ON DELETE CASCADE,
source_asset_id INTEGER REFERENCES image_studio_assets(id) ON DELETE SET NULL,
reference_asset_ids TEXT,
output_asset_id INTEGER REFERENCES image_studio_assets(id) ON DELETE SET NULL,
generation_source TEXT NOT NULL DEFAULT 'cmhub',
provider TEXT NOT NULL DEFAULT 'cmhub',
@@ -488,6 +489,7 @@ def init_db(path=None, conn=None) -> None:
_ensure_image_studio_project_draft_columns(database)
_ensure_image_studio_job_recovery_columns(database)
_ensure_image_studio_generation_round_columns(database)
_ensure_image_studio_job_reference_asset_ids_column(database)
def _ensure_batch_delete_columns(database):
@@ -612,6 +614,16 @@ def _ensure_image_studio_generation_round_columns(database):
"project_id, generation_round_key, generation_slot_index, id)"
)
def _ensure_image_studio_job_reference_asset_ids_column(database):
columns = {
row["name"] for row in database.execute("PRAGMA table_info(image_studio_jobs)").fetchall()
}
if "reference_asset_ids" not in columns:
database.execute(
"ALTER TABLE image_studio_jobs ADD COLUMN reference_asset_ids 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/tabs/product_suite.py
+7
View File
@@ -4704,6 +4704,7 @@ class ProductSuiteTab(QWidget):
return tuple(
(
int(spec.get("source_asset_id") or 0),
tuple(int(asset_id) for asset_id in (spec.get("reference_asset_ids") or [])),
str(spec.get("category") or spec.get("job_type") or ""),
int(spec.get("category_index") or 0),
int(spec.get("source_index") or 0),
@@ -5582,8 +5583,14 @@ class ProductSuiteTab(QWidget):
}:
self._status("当前图片无需重试", "warning")
return
try:
reference_asset_ids = image_studio.job_reference_asset_ids(job)
except Exception as exc:
self._message("读取重试图片失败", _user_error(exc))
return
spec = {
"source_asset_id": job.source_asset_id,
"reference_asset_ids": reference_asset_ids,
"job_type": job.job_type,
"prompt": job.prompt,
}
app/gui/workers.py
+1
View File
@@ -372,6 +372,7 @@ class ProductSuiteGenerateWorker(BaseWorker):
image_studio.create_job(
self.project_id,
source_asset_id=spec.get("source_asset_id"),
reference_asset_ids=spec.get("reference_asset_ids"),
job_type=spec.get("job_type") or "套图",
prompt=spec.get("prompt") or "",
generation_source="cmhub",
app/image_studio.py
+51 -3
View File
@@ -78,6 +78,7 @@ class ImageStudioJob:
id: int
project_id: int
source_asset_id: Optional[int]
reference_asset_ids: Optional[str]
output_asset_id: Optional[int]
generation_source: str
provider: str
@@ -344,6 +345,41 @@ def _ensure_assets_belong_to_project(database, project_id, asset_ids):
raise db.DbError("AI工场资产不属于当前项目")
def _parse_reference_asset_ids(value, source_asset_id=None):
if value is None or str(value).strip() == "":
return []
if isinstance(value, str):
try:
value = json.loads(value)
except (TypeError, ValueError, json.JSONDecodeError) as exc:
raise db.DbError("AI工场参考图快照格式无效") from exc
if not isinstance(value, (list, tuple)):
raise db.DbError("AI工场参考图快照必须是图片ID列表")
source_id = int(source_asset_id) if source_asset_id is not None else None
normalized = []
for asset_id in value:
try:
parsed = int(asset_id)
except (TypeError, ValueError) as exc:
raise db.DbError("AI工场参考图快照包含无效图片ID") from exc
if parsed <= 0:
raise db.DbError("AI工场参考图快照包含无效图片ID")
if source_id is not None and parsed == source_id:
raise db.DbError("AI工场参考图不能包含主图")
if parsed in normalized:
raise db.DbError("AI工场参考图不能重复")
normalized.append(parsed)
return normalized
def job_reference_asset_ids(job):
"""Return a validated, ordered reference asset snapshot for one job."""
return _parse_reference_asset_ids(
getattr(job, "reference_asset_ids", None),
getattr(job, "source_asset_id", None),
)
def get_project(project_id, path=None, conn=None, include_deleted=False):
clauses = ["id = ?"]
params = [int(project_id)]
@@ -1114,6 +1150,7 @@ def create_job(
project_id,
*,
source_asset_id=None,
reference_asset_ids=None,
job_type="main",
prompt="",
task_key=None,
@@ -1126,6 +1163,12 @@ def create_job(
):
now = _now()
task_key = str(task_key or _task_key(project_id))
reference_ids = _parse_reference_asset_ids(reference_asset_ids, source_asset_id)
reference_json = None if reference_asset_ids is None else json.dumps(
reference_ids,
ensure_ascii=True,
separators=(",", ":"),
)
if generation_round_key is not None:
generation_round_key = str(generation_round_key).strip()
if not generation_round_key:
@@ -1140,18 +1183,23 @@ def create_job(
with _connection(conn, path) as database:
try:
with database:
_ensure_assets_belong_to_project(database, project_id, [source_asset_id])
_ensure_assets_belong_to_project(
database,
project_id,
[source_asset_id, *reference_ids],
)
cursor = database.execute(
"""
INSERT INTO image_studio_jobs
(project_id, source_asset_id, generation_source, provider,
(project_id, source_asset_id, reference_asset_ids, generation_source, provider,
job_type, task_key, status, prompt, generation_round_key,
generation_slot_index, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?, 'pending', ?, ?, ?, ?, ?)
VALUES (?, ?, ?, ?, ?, ?, ?, 'pending', ?, ?, ?, ?, ?)
""",
(
int(project_id),
source_asset_id,
reference_json,
str(generation_source or "cmhub"),
str(provider or "cmhub"),
str(job_type),
app/image_studio_generation.py
+18
View File
@@ -270,6 +270,9 @@ def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, sho
return {"job": updated, "status": "failed", "error": "项目或源图不存在"}
try:
image_studio.update_job_status(job.id, "running", 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,
@@ -279,6 +282,7 @@ def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, sho
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)
@@ -343,6 +347,20 @@ def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, sho
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 _submit_or_resume_job(
job,
source_asset,
app/product_suite.py
+8
View File
@@ -347,6 +347,13 @@ def build_job_specs(source_assets, base_prompt, settings, item_id, *, template_t
if not assets:
return []
normalized = normalize_suite_settings(settings)
primary_asset_id = int(getattr(assets[0], "id", assets[0]))
reference_asset_ids = []
if not normalized["per_image_primary"]:
for asset in assets[1:8]:
asset_id = int(getattr(asset, "id", asset))
if asset_id != primary_asset_id and asset_id not in reference_asset_ids:
reference_asset_ids.append(asset_id)
specs = []
for category in category_order(normalized):
count = _count(normalized["categories"].get(category, 0))
@@ -358,6 +365,7 @@ def build_job_specs(source_assets, base_prompt, settings, item_id, *, template_t
specs.append(
{
"source_asset_id": int(getattr(asset, "id", asset)),
"reference_asset_ids": list(reference_asset_ids),
"job_type": str(category),
"category": str(category),
"category_index": category_index,