Spaces:

nenae18
/

leechard

Running

Add face refinement levels and salon orbit video

215580b verified 2 months ago

8.63 kB

"""비동기 변환 파이프라인 (Celery) — 멱등 + lease fencing.

API 서버는 enqueue 만 하고 즉시 응답한다. 이 태스크는:
  - 이미 종료 상태(SUCCEEDED/EXPIRED/DELETED)면 즉시 return (멱등)
  - QUEUED 인 task 만 lock 획득(LockContext 발급). 이미 lock 이 살아 있으면 no-op
  - 모든 상태 쓰기는 (locked_by, lock_token, lock_version) fencing 으로 보호
    → reaper/다른 worker 에게 lock 이 넘어간 뒤 늦게 살아난 worker 는 어떤 것도 못 쓴다
  - 결과는 attempt-scoped 경로에 저장하고, SUCCEEDED 전이에 성공한 worker 만 promotion

흐름:
  QUEUED → QC_CHECKING → QC_PASSED → AI_PROCESSING
        → POST_PROCESSING → QA_CHECKING → SUCCEEDED
실패 시 각 단계에서 *_FAILED 로 분기.
"""
from __future__ import annotations

import logging
import os

from app.db.session import SessionLocal
from app.domain.states import TERMINAL_STATES, TaskStatus
from app.services import artifact_repo, task_repo
from app.services.comfyui_client import FaceSwapRequest, get_ai_client
from app.services.photo_qc import get_photo_qc
from app.services.retention import run_retention_cleanup
from app.services.storage import get_storage
from app.services.task_repo import LockContext
from app.services.watermark import apply_ai_watermark
from app.workers.celery_app import celery_app

logger = logging.getLogger(__name__)

STALE = "STALE_ABORT"


def attempt_result_key(task_id: str, attempt_id: str) -> str:
    """attempt-scoped 결과 키. 덮어쓰기 대신 시도별 분리 → stale worker 충돌 방지."""
    return f"results/{task_id}/attempt-{attempt_id}.png"


@celery_app.task(name="process_photo", bind=True, max_retries=2)
def process_photo(self, task_id: str, dispatch_version: int = 0) -> str:
    db = SessionLocal()
    worker_id = f"{self.request.hostname or 'worker'}:{self.request.id or os.getpid()}"
    ctx: LockContext | None = None
    try:
        task = task_repo.get_task(db, task_id)
        if task is None:
            logger.warning("task not found: %s", task_id)
            return "NOT_FOUND"

        # 멱등: 이미 종료된 작업은 재처리하지 않는다.
        if task.status in TERMINAL_STATES:
            return f"SKIP_TERMINAL:{task.status.value}"

        # dispatch fencing: manual retry 등으로 세대가 올라간 old dispatch 는 no-op.
        if task.dispatch_version != dispatch_version:
            logger.info(
                "stale dispatch %s: payload dv=%s != task dv=%s",
                task_id, dispatch_version, task.dispatch_version,
            )
            return "SKIP_STALE_DISPATCH"

        # lock: QUEUED + 같은 dispatch_version 인 task 만 획득(fencing context 발급).
        ctx = task_repo.acquire_lock(
            db, task_id, worker_id, dispatch_version=dispatch_version
        )
        if ctx is None:
            logger.info("no-op (lock held or not QUEUED): %s", task_id)
            return "SKIP_LOCKED"
        db.refresh(task)

        storage = get_storage()
        qc = get_photo_qc()
        ai = get_ai_client()

        # 1) QUEUED -> QC_CHECKING
        if not task_repo.fenced_transition(
            db, ctx, TaskStatus.QC_CHECKING, expected=TaskStatus.QUEUED
        ):
            return STALE
        original = storage.get_bytes(task.original_key)
        # 폰 사진은 EXIF 방향 태그로 픽셀이 눕혀져 저장되는 경우가 많다. AI 엔진은
        # 내부에서 방향을 정규화(upright)해 결과를 내므로, 여기서 원본도 미리 정규화하지
        # 않으면 QC/AI/QA/출력이 서로 다른 방향(전치된 크기)을 보게 되어 정상 결과가
        # OUTPUT_SIZE_MISMATCH 로 잘못 걸린다. 적재 직후 한 번만 정규화해 모두 일치시킨다.
        from app.services.gemini_client import normalize_image_orientation_bytes
        original = normalize_image_orientation_bytes(original)

        # 2) 입력 QC
        in_qc = qc.check_input(original, task.original_mime or "image/jpeg")
        if not in_qc.passed:
            task_repo.fenced_transition(
                db, ctx, TaskStatus.QC_FAILED, expected=TaskStatus.QC_CHECKING,
                clear_lock=True, error_code=in_qc.code, error_detail=in_qc.detail,
                failure_code=in_qc.code, failure_stage="QC_CHECKING",
            )
            return TaskStatus.QC_FAILED.value
        if not task_repo.fenced_transition(
            db, ctx, TaskStatus.QC_PASSED, expected=TaskStatus.QC_CHECKING
        ):
            return STALE

        # 3) AI 변환 (얼굴 영역만)
        if not task_repo.fenced_transition(
            db, ctx, TaskStatus.AI_PROCESSING, expected=TaskStatus.QC_PASSED
        ):
            return STALE
        try:
            ai_res = ai.face_only_transform(
                FaceSwapRequest(
                    image_bytes=original,
                    mime=task.original_mime or "image/jpeg",
                    face_boxes=in_qc.face_boxes,
                    requested_gender=task.requested_gender,
                    beauty_level=task.beauty_level or "balanced",
                )
            )
        except Exception as exc:  # noqa: BLE001
            logger.exception("AI transform failed")
            code = getattr(exc, "code", "AI_ERROR")
            task_repo.fenced_transition(
                db, ctx, TaskStatus.AI_FAILED, expected=TaskStatus.AI_PROCESSING,
                clear_lock=True, error_code=code, error_detail=str(exc),
                failure_code=code, failure_message=str(exc),
                failure_stage="AI_PROCESSING",
            )
            return TaskStatus.AI_FAILED.value

        # 4) 후처리 (리사이즈/색보정 자리) — mock 통과
        if not task_repo.fenced_transition(
            db, ctx, TaskStatus.POST_PROCESSING, expected=TaskStatus.AI_PROCESSING
        ):
            return STALE
        result_bytes = ai_res.image_bytes

        # 5) 결과 QA — 얼굴 외 영역 보존 검증
        if not task_repo.fenced_transition(
            db, ctx, TaskStatus.QA_CHECKING, expected=TaskStatus.POST_PROCESSING
        ):
            return STALE
        out_qc = qc.check_output(original, result_bytes)
        if not out_qc.passed:
            task_repo.fenced_transition(
                db, ctx, TaskStatus.QA_FAILED, expected=TaskStatus.QA_CHECKING,
                clear_lock=True, error_code=out_qc.code, error_detail=out_qc.detail,
                failure_code=out_qc.code, failure_stage="QA_CHECKING",
            )
            return TaskStatus.QA_FAILED.value

        # 6) AI 고지 워터마크 + 메타데이터 (stub/real 동일). 공유/다운로드되는
        #    최종 이미지에만 적용하며, 얼굴/헤어를 가리지 않는 여백 모서리에 표시한다.
        result_bytes = apply_ai_watermark(result_bytes)

        # 7) attempt-scoped 경로에 저장 + artifact TEMP 기록
        #    (stale worker 도 자기 attempt 경로에만 쓴다 → 충돌 없음)
        result_key = attempt_result_key(task.task_id, ctx.current_attempt_id)
        storage.put_bytes(result_key, result_bytes, "image/png")
        result_url = storage.public_url(result_key)
        artifact_repo.record_temp(db, task.task_id, ctx.current_attempt_id, result_key)

        # 7) SUCCEEDED 전이 + result_key promotion (원자적, fencing).
        #    성공한 worker 만 artifact 를 PROMOTED 로 승격하고, stale 은 STALE 로 표시.
        if not task_repo.fenced_succeed(
            db, ctx, result_key=result_key, result_url=result_url
        ):
            logger.info("lost SUCCEEDED race (stale): %s", task_id)
            artifact_repo.mark_stale(db, task.task_id, ctx.current_attempt_id)
            return STALE
        artifact_repo.mark_promoted(db, task.task_id, ctx.current_attempt_id)
        return TaskStatus.SUCCEEDED.value

    except Exception as exc:  # noqa: BLE001
        logger.exception("pipeline crashed for %s", task_id)
        if ctx is not None:
            try:
                task_repo.release_lock(db, ctx)
            except Exception:  # noqa: BLE001
                pass
        raise self.retry(exc=exc, countdown=5)
    finally:
        db.close()


@celery_app.task(name="retention_cleanup")
def retention_cleanup() -> dict:
    db = SessionLocal()
    try:
        return run_retention_cleanup(db)
    finally:
        db.close()

	"""비동기 변환 파이프라인 (Celery) — 멱등 + lease fencing.

	API 서버는 enqueue 만 하고 즉시 응답한다. 이 태스크는:
	- 이미 종료 상태(SUCCEEDED/EXPIRED/DELETED)면 즉시 return (멱등)
	- QUEUED 인 task 만 lock 획득(LockContext 발급). 이미 lock 이 살아 있으면 no-op
	- 모든 상태 쓰기는 (locked_by, lock_token, lock_version) fencing 으로 보호
	→ reaper/다른 worker 에게 lock 이 넘어간 뒤 늦게 살아난 worker 는 어떤 것도 못 쓴다
	- 결과는 attempt-scoped 경로에 저장하고, SUCCEEDED 전이에 성공한 worker 만 promotion

	흐름:
	QUEUED → QC_CHECKING → QC_PASSED → AI_PROCESSING
	→ POST_PROCESSING → QA_CHECKING → SUCCEEDED
	실패 시 각 단계에서 *_FAILED 로 분기.
	"""
	from __future__ import annotations

	import logging
	import os

	from app.db.session import SessionLocal
	from app.domain.states import TERMINAL_STATES, TaskStatus
	from app.services import artifact_repo, task_repo
	from app.services.comfyui_client import FaceSwapRequest, get_ai_client
	from app.services.photo_qc import get_photo_qc
	from app.services.retention import run_retention_cleanup
	from app.services.storage import get_storage
	from app.services.task_repo import LockContext
	from app.services.watermark import apply_ai_watermark
	from app.workers.celery_app import celery_app

	logger = logging.getLogger(__name__)

	STALE = "STALE_ABORT"


	def attempt_result_key(task_id: str, attempt_id: str) -> str:
	"""attempt-scoped 결과 키. 덮어쓰기 대신 시도별 분리 → stale worker 충돌 방지."""
	return f"results/{task_id}/attempt-{attempt_id}.png"


	@celery_app.task(name="process_photo", bind=True, max_retries=2)
	def process_photo(self, task_id: str, dispatch_version: int = 0) -> str:
	db = SessionLocal()
	worker_id = f"{self.request.hostname or 'worker'}:{self.request.id or os.getpid()}"
	ctx: LockContext \| None = None
	try:
	task = task_repo.get_task(db, task_id)
	if task is None:
	logger.warning("task not found: %s", task_id)
	return "NOT_FOUND"

	# 멱등: 이미 종료된 작업은 재처리하지 않는다.
	if task.status in TERMINAL_STATES:
	return f"SKIP_TERMINAL:{task.status.value}"

	# dispatch fencing: manual retry 등으로 세대가 올라간 old dispatch 는 no-op.
	if task.dispatch_version != dispatch_version:
	logger.info(
	"stale dispatch %s: payload dv=%s != task dv=%s",
	task_id, dispatch_version, task.dispatch_version,
	)
	return "SKIP_STALE_DISPATCH"

	# lock: QUEUED + 같은 dispatch_version 인 task 만 획득(fencing context 발급).
	ctx = task_repo.acquire_lock(
	db, task_id, worker_id, dispatch_version=dispatch_version
	)
	if ctx is None:
	logger.info("no-op (lock held or not QUEUED): %s", task_id)
	return "SKIP_LOCKED"
	db.refresh(task)

	storage = get_storage()
	qc = get_photo_qc()
	ai = get_ai_client()

	# 1) QUEUED -> QC_CHECKING
	if not task_repo.fenced_transition(
	db, ctx, TaskStatus.QC_CHECKING, expected=TaskStatus.QUEUED
	):
	return STALE
	original = storage.get_bytes(task.original_key)
	# 폰 사진은 EXIF 방향 태그로 픽셀이 눕혀져 저장되는 경우가 많다. AI 엔진은
	# 내부에서 방향을 정규화(upright)해 결과를 내므로, 여기서 원본도 미리 정규화하지
	# 않으면 QC/AI/QA/출력이 서로 다른 방향(전치된 크기)을 보게 되어 정상 결과가
	# OUTPUT_SIZE_MISMATCH 로 잘못 걸린다. 적재 직후 한 번만 정규화해 모두 일치시킨다.
	from app.services.gemini_client import normalize_image_orientation_bytes
	original = normalize_image_orientation_bytes(original)

	# 2) 입력 QC
	in_qc = qc.check_input(original, task.original_mime or "image/jpeg")
	if not in_qc.passed:
	task_repo.fenced_transition(
	db, ctx, TaskStatus.QC_FAILED, expected=TaskStatus.QC_CHECKING,
	clear_lock=True, error_code=in_qc.code, error_detail=in_qc.detail,
	failure_code=in_qc.code, failure_stage="QC_CHECKING",
	)
	return TaskStatus.QC_FAILED.value
	if not task_repo.fenced_transition(
	db, ctx, TaskStatus.QC_PASSED, expected=TaskStatus.QC_CHECKING
	):
	return STALE

	# 3) AI 변환 (얼굴 영역만)
	if not task_repo.fenced_transition(
	db, ctx, TaskStatus.AI_PROCESSING, expected=TaskStatus.QC_PASSED
	):
	return STALE
	try:
	ai_res = ai.face_only_transform(
	FaceSwapRequest(
	image_bytes=original,
	mime=task.original_mime or "image/jpeg",
	face_boxes=in_qc.face_boxes,
	requested_gender=task.requested_gender,
	beauty_level=task.beauty_level or "balanced",
	)
	)
	except Exception as exc: # noqa: BLE001
	logger.exception("AI transform failed")
	code = getattr(exc, "code", "AI_ERROR")
	task_repo.fenced_transition(
	db, ctx, TaskStatus.AI_FAILED, expected=TaskStatus.AI_PROCESSING,
	clear_lock=True, error_code=code, error_detail=str(exc),
	failure_code=code, failure_message=str(exc),
	failure_stage="AI_PROCESSING",
	)
	return TaskStatus.AI_FAILED.value

	# 4) 후처리 (리사이즈/색보정 자리) — mock 통과
	if not task_repo.fenced_transition(
	db, ctx, TaskStatus.POST_PROCESSING, expected=TaskStatus.AI_PROCESSING
	):
	return STALE
	result_bytes = ai_res.image_bytes

	# 5) 결과 QA — 얼굴 외 영역 보존 검증
	if not task_repo.fenced_transition(
	db, ctx, TaskStatus.QA_CHECKING, expected=TaskStatus.POST_PROCESSING
	):
	return STALE
	out_qc = qc.check_output(original, result_bytes)
	if not out_qc.passed:
	task_repo.fenced_transition(
	db, ctx, TaskStatus.QA_FAILED, expected=TaskStatus.QA_CHECKING,
	clear_lock=True, error_code=out_qc.code, error_detail=out_qc.detail,
	failure_code=out_qc.code, failure_stage="QA_CHECKING",
	)
	return TaskStatus.QA_FAILED.value

	# 6) AI 고지 워터마크 + 메타데이터 (stub/real 동일). 공유/다운로드되는
	# 최종 이미지에만 적용하며, 얼굴/헤어를 가리지 않는 여백 모서리에 표시한다.
	result_bytes = apply_ai_watermark(result_bytes)

	# 7) attempt-scoped 경로에 저장 + artifact TEMP 기록
	# (stale worker 도 자기 attempt 경로에만 쓴다 → 충돌 없음)
	result_key = attempt_result_key(task.task_id, ctx.current_attempt_id)
	storage.put_bytes(result_key, result_bytes, "image/png")
	result_url = storage.public_url(result_key)
	artifact_repo.record_temp(db, task.task_id, ctx.current_attempt_id, result_key)

	# 7) SUCCEEDED 전이 + result_key promotion (원자적, fencing).
	# 성공한 worker 만 artifact 를 PROMOTED 로 승격하고, stale 은 STALE 로 표시.
	if not task_repo.fenced_succeed(
	db, ctx, result_key=result_key, result_url=result_url
	):
	logger.info("lost SUCCEEDED race (stale): %s", task_id)
	artifact_repo.mark_stale(db, task.task_id, ctx.current_attempt_id)
	return STALE
	artifact_repo.mark_promoted(db, task.task_id, ctx.current_attempt_id)
	return TaskStatus.SUCCEEDED.value

	except Exception as exc: # noqa: BLE001
	logger.exception("pipeline crashed for %s", task_id)
	if ctx is not None:
	try:
	task_repo.release_lock(db, ctx)
	except Exception: # noqa: BLE001
	pass
	raise self.retry(exc=exc, countdown=5)
	finally:
	db.close()


	@celery_app.task(name="retention_cleanup")
	def retention_cleanup() -> dict:
	db = SessionLocal()
	try:
	return run_retention_cleanup(db)
	finally:
	db.close()