455 lines
16 KiB
Python
455 lines
16 KiB
Python
"""
|
||
Story 主插图生成 Celery 任务。
|
||
|
||
从 story_image_intents 原子 claim intent，生成图片，写入 assets，更新 intent。
|
||
不读取正文占位符。
|
||
"""
|
||
|
||
import hashlib
|
||
import time
|
||
import uuid
|
||
from datetime import datetime, timedelta, timezone
|
||
|
||
from celery import shared_task
|
||
from PIL import Image
|
||
from sqlalchemy import and_, func, or_, select, update
|
||
|
||
from app.agents.image_prompt import get_image_prompt_orchestrator
|
||
from app.core.db import get_sync_db
|
||
from app.core.dependencies import get_image_generator
|
||
from app.core.logging import get_logger
|
||
from app.core.memoir_pipeline_progress import merge_fanout_item
|
||
from app.core.redis_lock import acquire_redis_lock, release_redis_lock
|
||
from app.features.asset.models import Asset
|
||
from app.features.memoir.asset_resolver import strip_asset_image_refs_from_markdown
|
||
from app.features.memoir.memoir_images.settings import MemoirImageSettings
|
||
from app.features.memoir.memoir_images.storage import TencentCosStorageService
|
||
from app.features.story.backfill import backfill_image_into_markdown
|
||
from app.features.story.models import Story, StoryImageIntent, StoryVersion
|
||
from app.ports.image_gen import TaskStatus
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
STORY_IMAGE_LOCK_TTL_SECONDS = 1800
|
||
STORY_IMAGE_CLAIM_TTL_SECONDS = 1800
|
||
|
||
|
||
def _enqueue_chapter_effects_after_image_backfill(story_id: str) -> None:
|
||
"""主图回填后标记关联章节 dirty，并经统一 post-commit 入口派发章节物化与 compaction。"""
|
||
try:
|
||
with get_sync_db() as session:
|
||
from app.features.memoir import repo as memoir_repo
|
||
|
||
story = session.get(Story, story_id)
|
||
if not story:
|
||
return
|
||
uid = str(story.user_id)
|
||
memoir_repo.mark_chapters_dirty_for_story_sync(session, story_id)
|
||
chapter_ids = memoir_repo.get_chapter_ids_linked_to_story_sync(
|
||
session, story_id
|
||
)
|
||
session.commit()
|
||
user_id = uid
|
||
except Exception as exc:
|
||
logger.warning(
|
||
"mark_chapters_dirty_for_story_sync failed story={}: {}", story_id, exc
|
||
)
|
||
return
|
||
from app.features.story.post_commit import enqueue_story_post_commit_effects
|
||
|
||
enqueue_story_post_commit_effects(
|
||
user_id=user_id,
|
||
story_ids=set(),
|
||
chapter_ids=set(chapter_ids),
|
||
trigger_source="story_image_backfill",
|
||
need_image=False,
|
||
need_recompose=True,
|
||
need_compaction=True,
|
||
)
|
||
|
||
|
||
def _build_story_image_cos_key(
|
||
user_id: str, story_id: str, intent_id: str, prompt: str
|
||
) -> str:
|
||
short_hash = hashlib.sha1(prompt.encode("utf-8")).hexdigest()[:10]
|
||
return f"stories/{user_id}/{story_id}/{intent_id}-{short_hash}.png"
|
||
|
||
|
||
def _normalize_image_bytes(image_bytes: bytes) -> bytes:
|
||
from io import BytesIO
|
||
|
||
with Image.open(BytesIO(image_bytes)) as image:
|
||
output = BytesIO()
|
||
if image.mode in {"RGBA", "LA"}:
|
||
normalized = image
|
||
elif image.mode == "P":
|
||
normalized = image.convert("RGBA")
|
||
else:
|
||
normalized = image.convert("RGB")
|
||
normalized.save(output, format="PNG")
|
||
return output.getvalue()
|
||
|
||
|
||
def _story_image_claimable_clause(now: datetime):
|
||
cutoff = now - timedelta(seconds=STORY_IMAGE_CLAIM_TTL_SECONDS)
|
||
return or_(
|
||
StoryImageIntent.status.in_(["pending", "failed"]),
|
||
and_(
|
||
StoryImageIntent.status == "processing",
|
||
or_(
|
||
StoryImageIntent.claimed_at.is_(None),
|
||
StoryImageIntent.claimed_at < cutoff,
|
||
),
|
||
),
|
||
)
|
||
|
||
|
||
def _claim_story_image_intent_sync(db, story_id: str, claim_token: str):
|
||
now = datetime.now(timezone.utc)
|
||
claimable = _story_image_claimable_clause(now)
|
||
candidate_id = db.execute(
|
||
select(StoryImageIntent.id)
|
||
.where(StoryImageIntent.story_id == story_id)
|
||
.where(StoryImageIntent.intent_role == "primary")
|
||
.where(claimable)
|
||
.order_by(
|
||
StoryImageIntent.updated_at.desc(), StoryImageIntent.created_at.desc()
|
||
)
|
||
.limit(1)
|
||
).scalar_one_or_none()
|
||
if not candidate_id:
|
||
return None
|
||
|
||
claimed = db.execute(
|
||
update(StoryImageIntent)
|
||
.where(StoryImageIntent.id == candidate_id)
|
||
.where(_story_image_claimable_clause(now))
|
||
.values(
|
||
status="processing",
|
||
claim_token=claim_token,
|
||
claimed_at=now,
|
||
updated_at=now,
|
||
error=None,
|
||
attempt_count=func.coalesce(StoryImageIntent.attempt_count, 0) + 1,
|
||
)
|
||
)
|
||
if (claimed.rowcount or 0) != 1:
|
||
db.rollback()
|
||
return None
|
||
|
||
row = (
|
||
db.execute(
|
||
select(StoryImageIntent, Story)
|
||
.join(Story, StoryImageIntent.story_id == Story.id)
|
||
.where(StoryImageIntent.id == candidate_id)
|
||
)
|
||
.unique()
|
||
.first()
|
||
)
|
||
db.commit()
|
||
return row
|
||
|
||
|
||
@shared_task(bind=True, max_retries=3, default_retry_delay=30)
|
||
def generate_story_image(self, story_id: str, memoir_correlation_id: str | None = None):
|
||
"""
|
||
为 story 生成主插图。
|
||
从 story_image_intents 原子认领 primary intent，生成后写入 assets 并更新 intent。
|
||
"""
|
||
celery_tid = str(self.request.id)
|
||
t0 = time.perf_counter()
|
||
logger.info(
|
||
"event=story_image_task_start story_id={} task_id={} msg=故事主图生成任务开始",
|
||
story_id,
|
||
celery_tid,
|
||
)
|
||
lock_key = f"lock:story-image:{story_id}"
|
||
lock_handle = acquire_redis_lock(lock_key, ttl_seconds=STORY_IMAGE_LOCK_TTL_SECONDS)
|
||
if lock_handle is None:
|
||
logger.debug("generate_story_image: story={}, reason=locked", story_id)
|
||
merge_fanout_item(
|
||
memoir_correlation_id,
|
||
list_name="story_images",
|
||
id_field="story_id",
|
||
item_id=story_id,
|
||
task_id=celery_tid,
|
||
status="locked",
|
||
)
|
||
return {"status": "locked"}
|
||
|
||
claim_token = uuid.uuid4().hex
|
||
intent = None
|
||
story = None
|
||
try:
|
||
with get_sync_db() as db:
|
||
row = _claim_story_image_intent_sync(db, story_id, claim_token)
|
||
if not row:
|
||
logger.debug(
|
||
"generate_story_image: story={}, reason=no_claimable_intent",
|
||
story_id,
|
||
)
|
||
merge_fanout_item(
|
||
memoir_correlation_id,
|
||
list_name="story_images",
|
||
id_field="story_id",
|
||
item_id=story_id,
|
||
task_id=celery_tid,
|
||
status="no_intent",
|
||
)
|
||
return {"status": "no_intent"}
|
||
|
||
intent, story = row
|
||
|
||
img_cfg = MemoirImageSettings.from_env()
|
||
min_body = img_cfg.story_image_min_body_chars
|
||
if min_body > 0:
|
||
plain = strip_asset_image_refs_from_markdown(
|
||
story.canonical_markdown or ""
|
||
).strip()
|
||
if len(plain) < min_body:
|
||
with get_sync_db() as db:
|
||
intent_db = db.get(StoryImageIntent, intent.id)
|
||
if intent_db and (intent_db.status or "").strip() == "processing":
|
||
intent_db.status = "skipped"
|
||
intent_db.error = f"body_below_min_chars:{len(plain)}"
|
||
intent_db.claim_token = None
|
||
intent_db.claimed_at = None
|
||
intent_db.updated_at = datetime.now(timezone.utc)
|
||
db.commit()
|
||
logger.info(
|
||
"generate_story_image: skipped body too short story={} len={} min={}",
|
||
story_id,
|
||
len(plain),
|
||
min_body,
|
||
)
|
||
merge_fanout_item(
|
||
memoir_correlation_id,
|
||
list_name="story_images",
|
||
id_field="story_id",
|
||
item_id=story_id,
|
||
task_id=celery_tid,
|
||
status="skipped_body_too_short",
|
||
)
|
||
return {"status": "skipped_body_too_short"}
|
||
|
||
merge_fanout_item(
|
||
memoir_correlation_id,
|
||
list_name="story_images",
|
||
id_field="story_id",
|
||
item_id=story_id,
|
||
task_id=celery_tid,
|
||
status="running",
|
||
)
|
||
|
||
image_generator = get_image_generator()
|
||
storage = TencentCosStorageService.from_env()
|
||
|
||
settings = img_cfg
|
||
orch = get_image_prompt_orchestrator()
|
||
intent_style = (getattr(intent, "style_profile", None) or "").strip() or None
|
||
prompt_out = orch.build_story_primary_prompt(
|
||
story_title=story.title or "",
|
||
story_stage=story.stage,
|
||
prompt_brief=intent.prompt_brief or "",
|
||
style_profile=intent_style,
|
||
)
|
||
prompt_final = prompt_out["prompt"]
|
||
style_for_image = (
|
||
(prompt_out.get("style") or "").strip()
|
||
or intent_style
|
||
or settings.default_style
|
||
)
|
||
result = image_generator.generate(
|
||
prompt_final,
|
||
prompt_out.get("size") or settings.default_size,
|
||
style_for_image,
|
||
)
|
||
if result.status != TaskStatus.COMPLETED or not result.image_url:
|
||
raise RuntimeError(result.error or "Image generation failed")
|
||
|
||
image_bytes = _normalize_image_bytes(
|
||
image_generator.download_image(result.image_url)
|
||
)
|
||
cos_key = _build_story_image_cos_key(
|
||
story.user_id, story_id, intent.id, prompt_final
|
||
)
|
||
url = storage.upload_bytes(image_bytes, cos_key, "image/png")
|
||
|
||
asset_id = str(uuid.uuid4())
|
||
with get_sync_db() as db:
|
||
intent_db = db.get(StoryImageIntent, intent.id)
|
||
if (
|
||
not intent_db
|
||
or (intent_db.status or "").strip() != "processing"
|
||
or (intent_db.claim_token or "").strip() != claim_token
|
||
):
|
||
logger.debug(
|
||
"generate_story_image: skip persist intent={} status={} claim={}",
|
||
intent.id,
|
||
getattr(intent_db, "status", None),
|
||
getattr(intent_db, "claim_token", None),
|
||
)
|
||
merge_fanout_item(
|
||
memoir_correlation_id,
|
||
list_name="story_images",
|
||
id_field="story_id",
|
||
item_id=story_id,
|
||
task_id=celery_tid,
|
||
status="superseded_or_cancelled",
|
||
)
|
||
return {"status": "superseded_or_cancelled"}
|
||
|
||
asset = Asset(
|
||
id=asset_id,
|
||
asset_type="story_image",
|
||
storage_key=cos_key,
|
||
url=url,
|
||
provider=settings.provider,
|
||
style_profile=style_for_image,
|
||
prompt_final=prompt_final,
|
||
status="completed",
|
||
)
|
||
db.add(asset)
|
||
db.flush()
|
||
|
||
story_db = db.get(Story, story_id)
|
||
target_vid = intent_db.story_version_id or story_db.current_version_id
|
||
current_vid = story_db.current_version_id
|
||
|
||
intent_db.asset_id = asset_id
|
||
intent_db.status = "completed"
|
||
intent_db.claim_token = None
|
||
intent_db.claimed_at = None
|
||
intent_db.error = None
|
||
intent_db.updated_at = datetime.now(timezone.utc)
|
||
db.flush()
|
||
|
||
# 仅当 intent 仍指向当前版本时回填正文，避免慢任务/重试把图插到新版本上
|
||
if not target_vid or target_vid != current_vid:
|
||
db.commit()
|
||
logger.debug(
|
||
"generate_story_image: stale intent skip backfill story={} "
|
||
"intent_ver={} current={} url={} asset={}",
|
||
story_id,
|
||
target_vid,
|
||
current_vid,
|
||
url,
|
||
asset_id,
|
||
)
|
||
merge_fanout_item(
|
||
memoir_correlation_id,
|
||
list_name="story_images",
|
||
id_field="story_id",
|
||
item_id=story_id,
|
||
task_id=celery_tid,
|
||
status="success_stale",
|
||
)
|
||
return {"status": "success_stale", "asset_id": asset_id}
|
||
|
||
ver = db.get(StoryVersion, target_vid)
|
||
if not ver:
|
||
db.commit()
|
||
merge_fanout_item(
|
||
memoir_correlation_id,
|
||
list_name="story_images",
|
||
id_field="story_id",
|
||
item_id=story_id,
|
||
task_id=celery_tid,
|
||
status="success_no_snapshot",
|
||
)
|
||
return {"status": "success_no_snapshot", "asset_id": asset_id}
|
||
|
||
base_md = strip_asset_image_refs_from_markdown(ver.markdown_snapshot or "")
|
||
alt_text = (getattr(intent_db, "prompt_brief", None) or "").strip()
|
||
if not alt_text:
|
||
alt_text = (getattr(intent_db, "caption", None) or "").strip()
|
||
backfilled_md = backfill_image_into_markdown(
|
||
base_md,
|
||
asset_id=asset_id,
|
||
image_alt=alt_text or "主插图",
|
||
)
|
||
max_stmt = select(func.max(StoryVersion.version_no)).where(
|
||
StoryVersion.story_id == story_id
|
||
)
|
||
max_no = db.execute(max_stmt).scalar()
|
||
version_no = (max_no or 0) + 1
|
||
new_ver = StoryVersion(
|
||
id=str(uuid.uuid4()),
|
||
story_id=story_id,
|
||
version_no=version_no,
|
||
markdown_snapshot=backfilled_md,
|
||
change_summary="主插图回填",
|
||
actor_type="system",
|
||
source_type="image_backfill",
|
||
parent_version_id=story_db.current_version_id,
|
||
)
|
||
db.add(new_ver)
|
||
db.flush()
|
||
story_db.current_version_id = new_ver.id
|
||
story_db.canonical_markdown = backfilled_md
|
||
|
||
db.commit()
|
||
|
||
_enqueue_chapter_effects_after_image_backfill(story_id)
|
||
|
||
ms = (time.perf_counter() - t0) * 1000
|
||
logger.info(
|
||
"event=story_image_task_done story_id={} asset_id={} duration_ms={:.1f} "
|
||
"msg=故事主图生成完成",
|
||
story_id,
|
||
asset_id,
|
||
ms,
|
||
)
|
||
logger.debug(
|
||
"generate_story_image: story={} asset={} url={} cos_key={} prompt_final={}",
|
||
story_id,
|
||
asset_id,
|
||
url,
|
||
cos_key,
|
||
prompt_final,
|
||
)
|
||
merge_fanout_item(
|
||
memoir_correlation_id,
|
||
list_name="story_images",
|
||
id_field="story_id",
|
||
item_id=story_id,
|
||
task_id=celery_tid,
|
||
status="success",
|
||
)
|
||
return {"status": "success", "asset_id": asset_id}
|
||
except Exception as exc:
|
||
if intent is not None:
|
||
with get_sync_db() as db:
|
||
intent_db = db.get(StoryImageIntent, intent.id)
|
||
if (
|
||
intent_db
|
||
and (intent_db.status or "").strip() != "completed"
|
||
and (intent_db.claim_token or "").strip() == claim_token
|
||
):
|
||
intent_db.status = "failed"
|
||
intent_db.claim_token = None
|
||
intent_db.claimed_at = None
|
||
intent_db.error = str(exc)
|
||
intent_db.updated_at = datetime.now(timezone.utc)
|
||
db.commit()
|
||
merge_fanout_item(
|
||
memoir_correlation_id,
|
||
list_name="story_images",
|
||
id_field="story_id",
|
||
item_id=story_id,
|
||
task_id=celery_tid,
|
||
status="failure",
|
||
extra={"error": str(exc)},
|
||
)
|
||
ms = (time.perf_counter() - t0) * 1000
|
||
logger.warning(
|
||
"event=story_image_task_failed story_id={} duration_ms={:.1f} error={} "
|
||
"msg=故事主图生成失败",
|
||
story_id,
|
||
ms,
|
||
exc,
|
||
)
|
||
raise self.retry(exc=exc) from exc
|
||
finally:
|
||
release_redis_lock(lock_handle)
|