""" generate.py — POST /generate Receives a job_id and mask JSON, determines the base image from the current step's result_image, creates a new step, and kicks off LaMa inpainting as a FastAPI BackgroundTask. Step creation logic ------------------- 1. Read job.current_step (call it N). 2. Use step N's result_image as the base image for the new step. 3. New step number = N + 1. 4. If step N+1 already exists (user went back then edited again): - Truncate the steps list to keep only 0..N. - Delete the orphaned step folders from disk to avoid clutter. This implements linear history: branching is not supported; new work after "Go Back" discards the forward history, like undo/redo in most editors. 5. Copy the base image to steps/{N+1}/original.jpg so each step folder is self-contained. 6. Write masks to steps/{N+1}/mask_hard.png and mask_soft.png. 7. Append a new StepRecord at step N+1 with status="processing". 8. Advance job.current_step to N+1. 9. Enqueue background inference. Engine access pattern --------------------- The LamaEngine is fetched from request.app.state while the Request is alive, then passed as a plain argument to run_inpaint_job() so it has no dependency on the Request object after the response is sent. """ from __future__ import annotations import shutil import traceback from pathlib import Path from typing import Literal import cv2 import numpy as np from fastapi import APIRouter, BackgroundTasks, HTTPException, Request from pydantic import BaseModel from app.config import STORAGE_DIR from app.utils import job_folder_name from app.jobs_store import ( get_job, update_job, add_step, update_step, truncate_steps_after, ) from app.services.mask_service import build_masks, save_mask from app.services.blend_service import blend, check_output_quality from app.services.inpaint_engine import BaseInpaintEngine from app.models.schemas import StepRecord router = APIRouter() # --------------------------------------------------------------------------- # Request schema # --------------------------------------------------------------------------- class GenerateRequest(BaseModel): job_id: str mask_source: Literal["manual", "ocr", "mixed"] = "manual" mask_json: dict # validated further inside mask_service.build_masks() # --------------------------------------------------------------------------- # Background worker # --------------------------------------------------------------------------- def run_inpaint_job( job_id: str, folder_name: str, step_number: int, image: np.ndarray, hard_mask: np.ndarray, soft_mask: np.ndarray, engine: BaseInpaintEngine, ) -> None: """Run LaMa inpainting and blend, then update the step's status. Called by FastAPI BackgroundTasks after the /generate response is sent. Updates the StepRecord to "completed" or "failed" — never raises. folder_name is the human-readable on-disk name (job_folder_name(job_id, created_at)), passed explicitly so the worker never needs to re-read the job from disk. """ step_dir = STORAGE_DIR / "watermark" / folder_name / "steps" / str(step_number) try: # Step 1 — run LaMa inpainting lama_output: np.ndarray = engine.inpaint(image, hard_mask) # Step 2 — composite output onto original using feathered mask result_image: np.ndarray = blend(image, lama_output, soft_mask) # Step 3 — quality check before saving passed, error_msg = check_output_quality(result_image, expected_shape=image.shape) if not passed: update_step(job_id, step_number, status="failed", error_message=error_msg) return # Step 4 — save result result_path = step_dir / "result.jpg" ok = cv2.imwrite(str(result_path), result_image) if not ok: update_step( job_id, step_number, status="failed", error_message="cv2.imwrite failed writing result.jpg.", ) return # Step 5 — mark completed update_step( job_id, step_number, status="completed", result_image=str(result_path), ) except Exception: update_step( job_id, step_number, status="failed", error_message=traceback.format_exc(), ) # --------------------------------------------------------------------------- # Route # --------------------------------------------------------------------------- @router.post("") async def generate( body: GenerateRequest, background_tasks: BackgroundTasks, request: Request, ) -> dict: """Validate mask JSON, create new step, enqueue background inpainting. Returns immediately with {"job_id": ..., "step_number": ..., "status": "processing"}. Inference has NOT started at the point this response is sent. """ job_id = body.job_id # --- Confirm job exists -------------------------------------------------- job = get_job(job_id) if job is None: raise HTTPException(status_code=404, detail=f"Job '{job_id}' not found.") # Compute the on-disk folder name once; reused for all path construction below. # job.created_at is immutable after creation so this always yields the same string. folder_name = job_folder_name(job.id, job.created_at) # --- Locate current step and base image ---------------------------------- current_step_num = job.current_step current_step = next( (s for s in job.steps if s.step_number == current_step_num), None ) if current_step is None or current_step.result_image is None: raise HTTPException( status_code=400, detail=f"Step {current_step_num} has no result image to build from.", ) base_path = Path(current_step.result_image) if not base_path.exists(): raise HTTPException( status_code=500, detail=f"Base image missing from storage: {base_path}", ) # --- Load base image ----------------------------------------------------- image: np.ndarray = cv2.imread(str(base_path)) if image is None: raise HTTPException( status_code=500, detail=f"cv2.imread could not decode: {base_path}", ) # --- Build masks --------------------------------------------------------- try: hard_mask, soft_mask = build_masks(body.mask_json) except ValueError as exc: raise HTTPException(status_code=422, detail=str(exc)) # --- Truncate any forward history (linear history: no branching) --------- new_step_num = current_step_num + 1 if any(s.step_number >= new_step_num for s in job.steps): removed = truncate_steps_after(job_id, keep_through=current_step_num) for s in removed: orphan = STORAGE_DIR / "watermark" / folder_name / "steps" / str(s["step_number"]) if orphan.exists(): shutil.rmtree(orphan) # --- Create step folder and copy base image as this step's original ------ step_dir = STORAGE_DIR / "watermark" / folder_name / "steps" / str(new_step_num) step_dir.mkdir(parents=True, exist_ok=True) step_original = step_dir / "original.jpg" shutil.copy2(str(base_path), str(step_original)) # --- Save masks ---------------------------------------------------------- hard_mask_path = step_dir / "mask_hard.png" soft_mask_path = step_dir / "mask_soft.png" save_mask(hard_mask, str(hard_mask_path)) save_mask(soft_mask, str(soft_mask_path)) # --- Append new step record (status=processing) -------------------------- new_step = StepRecord( step_number=new_step_num, original_image=str(step_original), mask_image=str(hard_mask_path), mask_source=body.mask_source, status="processing", ) add_step(job_id, new_step) # --- Advance current_step ------------------------------------------------ update_job(job_id, current_step=new_step_num) # --- Enqueue background inference ---------------------------------------- engine: BaseInpaintEngine = request.app.state.inpaint_engine background_tasks.add_task( run_inpaint_job, job_id = job_id, folder_name = folder_name, step_number = new_step_num, image = image, hard_mask = hard_mask, soft_mask = soft_mask, engine = engine, ) return {"job_id": job_id, "step_number": new_step_num, "status": "processing"}