566 lines
18 KiB
Python
566 lines
18 KiB
Python
"""Data service for the AI image studio."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import uuid
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from typing import Iterable, Optional
|
|
|
|
from . import appconfig, db
|
|
from .config import make_slug
|
|
|
|
|
|
PROJECT_STATUS_ACTIVE = "active"
|
|
ASSET_STATUS_AVAILABLE = "available"
|
|
ASSET_STATUS_MISSING = "missing"
|
|
ASSET_STATUSES = {ASSET_STATUS_AVAILABLE, ASSET_STATUS_MISSING}
|
|
JOB_STATUSES = {"pending", "submitted", "running", "succeeded", "failed", "expired", "cancelled"}
|
|
JOB_RESUMABLE_STATUSES = {"submitted", "running"}
|
|
SELECTION_TYPES = {"main", "detail"}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ImageStudioProject:
|
|
id: int
|
|
account_alias: str
|
|
account_slug: str
|
|
account_name: Optional[str]
|
|
item_id: str
|
|
target_main_count: int
|
|
target_detail_count: int
|
|
draft_prompt: Optional[str]
|
|
status: str
|
|
created_at: str
|
|
updated_at: str
|
|
deleted_at: Optional[str]
|
|
deleted_reason: Optional[str]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ImageStudioAsset:
|
|
id: int
|
|
project_id: int
|
|
kind: str
|
|
remote_url: Optional[str]
|
|
local_path: Optional[str]
|
|
aspect_ratio: Optional[str]
|
|
parent_asset_id: Optional[int]
|
|
prompt: Optional[str]
|
|
status: str
|
|
source_order: int
|
|
created_at: str
|
|
updated_at: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ImageStudioJob:
|
|
id: int
|
|
project_id: int
|
|
source_asset_id: Optional[int]
|
|
output_asset_id: Optional[int]
|
|
generation_source: str
|
|
provider: str
|
|
job_type: str
|
|
task_key: str
|
|
task_id: Optional[str]
|
|
status: str
|
|
prompt: Optional[str]
|
|
error: Optional[str]
|
|
attempts: int
|
|
points_cost: Optional[int]
|
|
points_balance: Optional[int]
|
|
call_id: Optional[str]
|
|
created_at: str
|
|
updated_at: str
|
|
submitted_at: Optional[str]
|
|
finished_at: Optional[str]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ImageStudioSelection:
|
|
id: int
|
|
project_id: int
|
|
selection_type: str
|
|
position: int
|
|
asset_id: int
|
|
created_at: str
|
|
updated_at: str
|
|
|
|
|
|
def _now() -> str:
|
|
return datetime.now().isoformat(timespec="seconds")
|
|
|
|
|
|
@contextmanager
|
|
def _connection(conn=None, path=None):
|
|
with db._connection(conn, path) as database:
|
|
yield database
|
|
|
|
|
|
def _row_to_dataclass(row, cls):
|
|
return None if row is None else cls(**dict(row))
|
|
|
|
|
|
def _fetch_one(conn, sql, params, cls):
|
|
return _row_to_dataclass(conn.execute(sql, params).fetchone(), cls)
|
|
|
|
|
|
def _fetch_all(conn, sql, params, cls):
|
|
return [_row_to_dataclass(row, cls) for row in conn.execute(sql, params).fetchall()]
|
|
|
|
|
|
def _get(obj, name, default=None):
|
|
if obj is None:
|
|
return default
|
|
if isinstance(obj, dict):
|
|
return obj.get(name, default)
|
|
return getattr(obj, name, default)
|
|
|
|
|
|
def _safe_component(value, default):
|
|
text = str(value or "").strip()
|
|
if not text:
|
|
text = str(default)
|
|
safe = "".join(ch if ch.isalnum() or ch in "_-" else "_" for ch in text).strip("_")
|
|
return safe or str(default)
|
|
|
|
|
|
def _account_fields(account=None, account_alias=None, account_name=None, account_slug=None):
|
|
alias = str(account_alias or _get(account, "alias") or "").strip()
|
|
if not alias:
|
|
raise db.DbError("AI工场项目缺少账号别名")
|
|
name = str(account_name or _get(account, "account_name") or "").strip() or None
|
|
slug = str(account_slug or _get(account, "slug") or "").strip() or make_slug(alias)
|
|
return alias, name, _safe_component(slug, "account")
|
|
|
|
|
|
def _task_key(project_id):
|
|
return f"image-studio-{int(project_id)}-{uuid.uuid4().hex}"
|
|
|
|
|
|
def _ensure_assets_belong_to_project(database, project_id, asset_ids):
|
|
ids = [int(asset_id) for asset_id in asset_ids if asset_id is not None]
|
|
if not ids:
|
|
return
|
|
placeholders = ",".join("?" for _ in ids)
|
|
rows = database.execute(
|
|
f"SELECT id FROM image_studio_assets WHERE project_id = ? AND id IN ({placeholders})",
|
|
[int(project_id), *ids],
|
|
).fetchall()
|
|
found = {int(row["id"]) for row in rows}
|
|
missing = sorted(set(ids) - found)
|
|
if missing:
|
|
raise db.DbError("AI工场资产不属于当前项目")
|
|
|
|
|
|
def get_project(project_id, path=None, conn=None, include_deleted=False):
|
|
clauses = ["id = ?"]
|
|
params = [int(project_id)]
|
|
if not include_deleted:
|
|
clauses.append("deleted_at IS NULL")
|
|
sql = "SELECT * FROM image_studio_projects WHERE " + " AND ".join(clauses)
|
|
with _connection(conn, path) as database:
|
|
return _fetch_one(database, sql, params, ImageStudioProject)
|
|
|
|
|
|
def get_project_by_account_item(account_alias, item_id, path=None, conn=None, include_deleted=False):
|
|
clauses = ["account_alias = ?", "item_id = ?"]
|
|
params = [str(account_alias).strip(), str(item_id).strip()]
|
|
if not include_deleted:
|
|
clauses.append("deleted_at IS NULL")
|
|
sql = "SELECT * FROM image_studio_projects WHERE " + " AND ".join(clauses)
|
|
with _connection(conn, path) as database:
|
|
return _fetch_one(database, sql, params, ImageStudioProject)
|
|
|
|
|
|
def create_or_get_project(
|
|
account=None,
|
|
item_id=None,
|
|
*,
|
|
account_alias=None,
|
|
account_name=None,
|
|
account_slug=None,
|
|
target_main_count=9,
|
|
target_detail_count=12,
|
|
draft_prompt="",
|
|
path=None,
|
|
conn=None,
|
|
):
|
|
alias, name, slug = _account_fields(
|
|
account,
|
|
account_alias=account_alias,
|
|
account_name=account_name,
|
|
account_slug=account_slug,
|
|
)
|
|
item = str(item_id or _get(account, "item_id") or "").strip()
|
|
if not item:
|
|
raise db.DbError("AI工场项目缺少商品ID")
|
|
now = _now()
|
|
with _connection(conn, path) as database:
|
|
existing = get_project_by_account_item(alias, item, conn=database, include_deleted=True)
|
|
if existing is not None:
|
|
if existing.deleted_at is not None:
|
|
with database:
|
|
database.execute(
|
|
"""
|
|
UPDATE image_studio_projects
|
|
SET deleted_at = NULL,
|
|
deleted_reason = NULL,
|
|
status = ?,
|
|
updated_at = ?
|
|
WHERE id = ?
|
|
""",
|
|
(PROJECT_STATUS_ACTIVE, now, existing.id),
|
|
)
|
|
return get_project(existing.id, conn=database)
|
|
return existing
|
|
with database:
|
|
cursor = database.execute(
|
|
"""
|
|
INSERT INTO image_studio_projects
|
|
(account_alias, account_slug, account_name, item_id,
|
|
target_main_count, target_detail_count, draft_prompt,
|
|
status, created_at, updated_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
alias,
|
|
slug,
|
|
name,
|
|
item,
|
|
int(target_main_count),
|
|
int(target_detail_count),
|
|
str(draft_prompt or ""),
|
|
PROJECT_STATUS_ACTIVE,
|
|
now,
|
|
now,
|
|
),
|
|
)
|
|
project_id = cursor.lastrowid
|
|
return get_project(project_id, conn=database)
|
|
|
|
|
|
def list_projects(path=None, conn=None, include_deleted=False):
|
|
sql = "SELECT * FROM image_studio_projects"
|
|
if not include_deleted:
|
|
sql += " WHERE deleted_at IS NULL"
|
|
sql += " ORDER BY updated_at DESC, id DESC"
|
|
with _connection(conn, path) as database:
|
|
return _fetch_all(database, sql, (), ImageStudioProject)
|
|
|
|
|
|
def update_project_prompt(project_id, draft_prompt, path=None, conn=None):
|
|
with _connection(conn, path) as database:
|
|
with database:
|
|
database.execute(
|
|
"""
|
|
UPDATE image_studio_projects
|
|
SET draft_prompt = ?, updated_at = ?
|
|
WHERE id = ? AND deleted_at IS NULL
|
|
""",
|
|
(str(draft_prompt or ""), _now(), int(project_id)),
|
|
)
|
|
return get_project(project_id, conn=database)
|
|
|
|
|
|
def soft_delete_project(project_id, reason="", path=None, conn=None):
|
|
now = _now()
|
|
with _connection(conn, path) as database:
|
|
with database:
|
|
database.execute(
|
|
"""
|
|
UPDATE image_studio_projects
|
|
SET deleted_at = ?, deleted_reason = ?, updated_at = ?
|
|
WHERE id = ? AND deleted_at IS NULL
|
|
""",
|
|
(now, str(reason or ""), now, int(project_id)),
|
|
)
|
|
return get_project(project_id, conn=database, include_deleted=True)
|
|
|
|
|
|
def project_image_dirs(image_root, project):
|
|
root = os.path.abspath(
|
|
os.path.join(
|
|
str(image_root or "images"),
|
|
"pool",
|
|
_safe_component(_get(project, "account_slug"), "account"),
|
|
_safe_component(_get(project, "item_id"), "item"),
|
|
)
|
|
)
|
|
return {
|
|
"root": root,
|
|
"originals": os.path.join(root, "originals"),
|
|
"generated": os.path.join(root, "generated"),
|
|
"exports": os.path.join(root, "exports"),
|
|
}
|
|
|
|
|
|
def default_project_image_dirs(project, config=None):
|
|
return project_image_dirs(appconfig.image_dir(config), project)
|
|
|
|
|
|
def add_asset(
|
|
project_id,
|
|
kind,
|
|
*,
|
|
remote_url=None,
|
|
local_path=None,
|
|
aspect_ratio=None,
|
|
parent_asset_id=None,
|
|
prompt=None,
|
|
status=ASSET_STATUS_AVAILABLE,
|
|
source_order=0,
|
|
path=None,
|
|
conn=None,
|
|
):
|
|
now = _now()
|
|
abs_local_path = os.path.abspath(local_path) if local_path else None
|
|
with _connection(conn, path) as database:
|
|
with database:
|
|
cursor = database.execute(
|
|
"""
|
|
INSERT INTO image_studio_assets
|
|
(project_id, kind, remote_url, local_path, aspect_ratio,
|
|
parent_asset_id, prompt, status, source_order, created_at, updated_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
int(project_id),
|
|
str(kind),
|
|
remote_url,
|
|
abs_local_path,
|
|
aspect_ratio,
|
|
parent_asset_id,
|
|
prompt,
|
|
str(status or ASSET_STATUS_AVAILABLE),
|
|
int(source_order or 0),
|
|
now,
|
|
now,
|
|
),
|
|
)
|
|
asset_id = cursor.lastrowid
|
|
return get_asset(asset_id, conn=database)
|
|
|
|
|
|
def get_asset(asset_id, path=None, conn=None):
|
|
with _connection(conn, path) as database:
|
|
return _fetch_one(
|
|
database,
|
|
"SELECT * FROM image_studio_assets WHERE id = ?",
|
|
(int(asset_id),),
|
|
ImageStudioAsset,
|
|
)
|
|
|
|
|
|
def list_assets(project_id, kind=None, include_missing=True, path=None, conn=None):
|
|
clauses = ["project_id = ?"]
|
|
params = [int(project_id)]
|
|
if kind is not None:
|
|
clauses.append("kind = ?")
|
|
params.append(str(kind))
|
|
if not include_missing:
|
|
clauses.append("status != ?")
|
|
params.append(ASSET_STATUS_MISSING)
|
|
sql = "SELECT * FROM image_studio_assets WHERE " + " AND ".join(clauses)
|
|
sql += " ORDER BY source_order, id"
|
|
with _connection(conn, path) as database:
|
|
return _fetch_all(database, sql, params, ImageStudioAsset)
|
|
|
|
|
|
def mark_asset_status(asset_id, status, path=None, conn=None):
|
|
if str(status) not in ASSET_STATUSES:
|
|
raise db.DbError("AI工场资产状态必须是 available 或 missing")
|
|
with _connection(conn, path) as database:
|
|
with database:
|
|
database.execute(
|
|
"""
|
|
UPDATE image_studio_assets
|
|
SET status = ?, updated_at = ?
|
|
WHERE id = ?
|
|
""",
|
|
(str(status), _now(), int(asset_id)),
|
|
)
|
|
return get_asset(asset_id, conn=database)
|
|
|
|
|
|
def create_job(
|
|
project_id,
|
|
*,
|
|
source_asset_id=None,
|
|
job_type="main",
|
|
prompt="",
|
|
task_key=None,
|
|
generation_source="cmhub",
|
|
provider="cmhub",
|
|
path=None,
|
|
conn=None,
|
|
):
|
|
now = _now()
|
|
task_key = str(task_key or _task_key(project_id))
|
|
with _connection(conn, path) as database:
|
|
try:
|
|
with database:
|
|
_ensure_assets_belong_to_project(database, project_id, [source_asset_id])
|
|
cursor = database.execute(
|
|
"""
|
|
INSERT INTO image_studio_jobs
|
|
(project_id, source_asset_id, generation_source, provider,
|
|
job_type, task_key, status, prompt, created_at, updated_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, 'pending', ?, ?, ?)
|
|
""",
|
|
(
|
|
int(project_id),
|
|
source_asset_id,
|
|
str(generation_source or "cmhub"),
|
|
str(provider or "cmhub"),
|
|
str(job_type),
|
|
task_key,
|
|
str(prompt or ""),
|
|
now,
|
|
now,
|
|
),
|
|
)
|
|
job_id = cursor.lastrowid
|
|
except Exception as exc:
|
|
raise db.DbError(f"创建AI工场生图任务失败: {exc}") from exc
|
|
return get_job(job_id, conn=database)
|
|
|
|
|
|
def get_job(job_id, path=None, conn=None):
|
|
with _connection(conn, path) as database:
|
|
return _fetch_one(
|
|
database,
|
|
"SELECT * FROM image_studio_jobs WHERE id = ?",
|
|
(int(job_id),),
|
|
ImageStudioJob,
|
|
)
|
|
|
|
|
|
def set_job_submitted(job_id, task_id, *, call_id=None, points_cost=None, points_balance=None, path=None, conn=None):
|
|
now = _now()
|
|
with _connection(conn, path) as database:
|
|
with database:
|
|
database.execute(
|
|
"""
|
|
UPDATE image_studio_jobs
|
|
SET task_id = ?,
|
|
status = 'submitted',
|
|
call_id = ?,
|
|
points_cost = ?,
|
|
points_balance = ?,
|
|
submitted_at = ?,
|
|
updated_at = ?
|
|
WHERE id = ?
|
|
""",
|
|
(str(task_id), call_id, points_cost, points_balance, now, now, int(job_id)),
|
|
)
|
|
return get_job(job_id, conn=database)
|
|
|
|
|
|
def update_job_status(
|
|
job_id,
|
|
status,
|
|
*,
|
|
error=None,
|
|
output_asset_id=None,
|
|
points_balance=None,
|
|
increment_attempts=False,
|
|
path=None,
|
|
conn=None,
|
|
):
|
|
if str(status) not in JOB_STATUSES:
|
|
raise db.DbError("AI工场任务状态无效")
|
|
now = _now()
|
|
terminal = str(status) in {"succeeded", "failed", "expired", "cancelled"}
|
|
with _connection(conn, path) as database:
|
|
with database:
|
|
if output_asset_id is not None:
|
|
job = get_job(job_id, conn=database)
|
|
if job is None:
|
|
raise db.DbError("AI工场任务不存在")
|
|
_ensure_assets_belong_to_project(database, job.project_id, [output_asset_id])
|
|
database.execute(
|
|
"""
|
|
UPDATE image_studio_jobs
|
|
SET status = ?,
|
|
error = ?,
|
|
output_asset_id = COALESCE(?, output_asset_id),
|
|
points_balance = COALESCE(?, points_balance),
|
|
attempts = attempts + ?,
|
|
finished_at = CASE WHEN ? THEN ? ELSE finished_at END,
|
|
updated_at = ?
|
|
WHERE id = ?
|
|
""",
|
|
(
|
|
str(status),
|
|
error,
|
|
output_asset_id,
|
|
points_balance,
|
|
1 if increment_attempts else 0,
|
|
1 if terminal else 0,
|
|
now,
|
|
now,
|
|
int(job_id),
|
|
),
|
|
)
|
|
return get_job(job_id, conn=database)
|
|
|
|
|
|
def list_resumable_jobs(path=None, conn=None, project_id=None):
|
|
clauses = ["status IN (?, ?)", "task_id IS NOT NULL"]
|
|
params = ["submitted", "running"]
|
|
if project_id is not None:
|
|
clauses.append("project_id = ?")
|
|
params.append(int(project_id))
|
|
sql = "SELECT * FROM image_studio_jobs WHERE " + " AND ".join(clauses)
|
|
sql += " ORDER BY updated_at, id"
|
|
with _connection(conn, path) as database:
|
|
return _fetch_all(database, sql, params, ImageStudioJob)
|
|
|
|
|
|
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:
|
|
raise db.DbError("终选类型必须是 main 或 detail")
|
|
ids = [int(asset_id) for asset_id in asset_ids]
|
|
if len(ids) != len(set(ids)):
|
|
raise db.DbError("同一资产不能重复加入同一类终选")
|
|
now = _now()
|
|
with _connection(conn, path) as database:
|
|
with database:
|
|
_ensure_assets_belong_to_project(database, project_id, ids)
|
|
database.execute(
|
|
"""
|
|
DELETE FROM image_studio_selections
|
|
WHERE project_id = ? AND selection_type = ?
|
|
""",
|
|
(int(project_id), selection),
|
|
)
|
|
database.executemany(
|
|
"""
|
|
INSERT INTO image_studio_selections
|
|
(project_id, selection_type, position, asset_id, created_at, updated_at)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
""",
|
|
[
|
|
(int(project_id), selection, index + 1, asset_id, now, now)
|
|
for index, asset_id in enumerate(ids)
|
|
],
|
|
)
|
|
return list_selections(project_id, selection, conn=database)
|
|
|
|
|
|
def list_selections(project_id, selection_type=None, path=None, conn=None):
|
|
clauses = ["project_id = ?"]
|
|
params = [int(project_id)]
|
|
if selection_type is not None:
|
|
clauses.append("selection_type = ?")
|
|
params.append(str(selection_type))
|
|
sql = "SELECT * FROM image_studio_selections WHERE " + " AND ".join(clauses)
|
|
sql += " ORDER BY selection_type, position"
|
|
with _connection(conn, path) as database:
|
|
return _fetch_all(database, sql, params, ImageStudioSelection)
|