Spaces:
Sleeping
Sleeping
| """ | |
| 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 | |
| # --------------------------------------------------------------------------- | |
| 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"} | |