paintplus: vendor the app source and rename from EditmaskwithAI
Bring the full EditmaskwithAI application into the repo under paintplus/ (429 files) so the service is self-contained — the installer copies the vendored source to ~/docker/paintplus/src instead of cloning at runtime. Rename to PaintPlus (service + branding; app logic untouched): - services/editmaskwithai.sh -> services/paintplus.sh (register_service paintplus, install_paintplus, ~/docker/paintplus, Caddy paintplus:8000, Authelia option preserved) - container names -> paintplus across docker-compose*.yml; dev network -> paintplus-network - browser <title> -> "PaintPlus - AI Image Editor"; README heading -> PaintPlus with upstream provenance note - README utilities table: editmaskwithai -> paintplus Backend/frontend code (help strings referencing the old container name, the ai_photo_edit.db filename) is intentionally left as-is to avoid touching application logic. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Nb2vJ8W7bHKx1JXVvpCraH
This commit is contained in:
Claude
@@ -0,0 +1,989 @@
|
||||
"""
|
||||
AI tools router — LaMa inpaint, background removal, remote generation, config.
|
||||
All endpoints are under /api prefix.
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional
|
||||
import base64
|
||||
import asyncio
|
||||
import json
|
||||
from io import BytesIO
|
||||
|
||||
from app.services.local_inpaint import (
|
||||
lama_inpaint, opencv_inpaint, lama_available, gpu_available, rembg_available,
|
||||
)
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["ai-tools"])
|
||||
|
||||
|
||||
# ─── Request / response models ───────────────────────────────────────────────
|
||||
|
||||
class EraseRequest(BaseModel):
|
||||
image: str # base64
|
||||
mask: str # base64
|
||||
|
||||
|
||||
class InpaintRemoteRequest(BaseModel):
|
||||
image: str
|
||||
mask: str
|
||||
prompt: str
|
||||
negative_prompt: Optional[str] = ""
|
||||
steps: Optional[int] = 30
|
||||
cfg_scale: Optional[float] = 7.5
|
||||
model: Optional[str] = None
|
||||
|
||||
|
||||
class Txt2ImgRequest(BaseModel):
|
||||
prompt: str
|
||||
width: Optional[int] = 1024
|
||||
height: Optional[int] = 1024
|
||||
negative_prompt: Optional[str] = ""
|
||||
steps: Optional[int] = 30
|
||||
cfg_scale: Optional[float] = 7.5
|
||||
model: Optional[str] = None
|
||||
seed: Optional[int] = 0
|
||||
|
||||
|
||||
class Img2ImgRequest(BaseModel):
|
||||
image: str
|
||||
prompt: str
|
||||
strength: Optional[float] = 0.75
|
||||
negative_prompt: Optional[str] = ""
|
||||
steps: Optional[int] = 30
|
||||
cfg_scale: Optional[float] = 7.5
|
||||
model: Optional[str] = None
|
||||
|
||||
|
||||
class OutpaintRequest(BaseModel):
|
||||
image: str
|
||||
direction: str # left | right | top | bottom
|
||||
size: Optional[int] = 256
|
||||
prompt: Optional[str] = ""
|
||||
|
||||
|
||||
class BgRemoveRequest(BaseModel):
|
||||
image: str
|
||||
|
||||
|
||||
# ─── Helpers ─────────────────────────────────────────────────────────────────
|
||||
|
||||
def _decode(b64: str) -> bytes:
|
||||
return base64.b64decode(b64)
|
||||
|
||||
|
||||
def _encode(data: bytes) -> str:
|
||||
return base64.b64encode(data).decode()
|
||||
|
||||
|
||||
def _require_remote(operation: str = None):
|
||||
from app.services.remote_provider import get_remote_provider
|
||||
from app.config import settings
|
||||
provider = get_remote_provider(operation)
|
||||
if provider is None:
|
||||
if (settings.ai_provider or "").lower() == "local_gpu":
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail=(
|
||||
"local_gpu provider failed to load — diffusers may be incompatible with "
|
||||
"the installed PyTorch version. Check container logs for details. "
|
||||
"If you see 'torch has no attribute xpu', rebuild the container from the "
|
||||
"correct branch so the pinned diffusers<0.29.0 is installed."
|
||||
)
|
||||
)
|
||||
op_hint = f"AI_PROVIDER_{operation.upper()} or " if operation else ""
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail=f"No remote AI provider configured for '{operation or 'default'}'. "
|
||||
f"Set {op_hint}AI_PROVIDER in .env (openai / invokeai / comfyui)."
|
||||
)
|
||||
return provider
|
||||
|
||||
|
||||
# ─── Local inpaint endpoints ─────────────────────────────────────────────────
|
||||
|
||||
@router.post("/erase")
|
||||
async def erase(req: EraseRequest):
|
||||
"""
|
||||
Magic eraser: remove object / fill region using LaMa (local, no API key needed).
|
||||
Falls back to OpenCV if LaMa not installed.
|
||||
"""
|
||||
try:
|
||||
image_bytes = _decode(req.image)
|
||||
mask_bytes = _decode(req.mask)
|
||||
|
||||
if lama_available():
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, lama_inpaint, image_bytes, mask_bytes
|
||||
)
|
||||
method = "lama"
|
||||
else:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, opencv_inpaint, image_bytes, mask_bytes
|
||||
)
|
||||
method = "opencv"
|
||||
|
||||
return {"result": _encode(result), "method": method}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/inpaint/lama")
|
||||
async def inpaint_lama(req: EraseRequest):
|
||||
"""LaMa structural inpainting."""
|
||||
if not lama_available():
|
||||
raise HTTPException(status_code=503, detail="simple-lama-inpainting not installed.")
|
||||
try:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, lama_inpaint, _decode(req.image), _decode(req.mask)
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/inpaint/fast")
|
||||
async def inpaint_fast(req: EraseRequest):
|
||||
"""OpenCV fast inpainting (CPU, milliseconds)."""
|
||||
try:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, opencv_inpaint, _decode(req.image), _decode(req.mask)
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/background/remove")
|
||||
async def background_remove(req: BgRemoveRequest):
|
||||
"""Remove background — rembg if available, else U2Net."""
|
||||
try:
|
||||
image_bytes = _decode(req.image)
|
||||
|
||||
# Try rembg first
|
||||
if rembg_available():
|
||||
from app.services.local_inpaint import remove_background_rembg
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, remove_background_rembg, image_bytes
|
||||
)
|
||||
return {"result": _encode(result), "method": "rembg"}
|
||||
|
||||
# Fall back to U2Net (existing implementation)
|
||||
from PIL import Image
|
||||
from io import BytesIO as _BytesIO
|
||||
img = Image.open(_BytesIO(image_bytes)).convert("RGB")
|
||||
from app.routers.tools import _remove_background_u2net
|
||||
result = await _remove_background_u2net(img)
|
||||
return {"result": _encode(result), "method": "u2net"}
|
||||
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# ─── Remote provider endpoints ───────────────────────────────────────────────
|
||||
|
||||
@router.post("/inpaint/remote")
|
||||
async def inpaint_remote(req: InpaintRemoteRequest):
|
||||
"""Inpaint via configured remote provider (InvokeAI / ComfyUI / OpenAI)."""
|
||||
provider = _require_remote("inpaint")
|
||||
try:
|
||||
params = {
|
||||
"negative_prompt": req.negative_prompt or "",
|
||||
"steps": req.steps,
|
||||
"cfg_scale": req.cfg_scale,
|
||||
}
|
||||
if req.model:
|
||||
params["model"] = req.model
|
||||
result = await provider.inpaint(_decode(req.image), _decode(req.mask), req.prompt, params)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/generate/progress")
|
||||
async def generation_progress_stream():
|
||||
"""
|
||||
SSE stream of local GPU pipeline inference progress.
|
||||
Events are JSON arrays of pipeline state objects, emitted every 200 ms.
|
||||
Each object: {pipeline, state, step, total_steps, progress, message, model_id, …}
|
||||
Clients open this with EventSource before firing a generation POST,
|
||||
then close it when the POST resolves.
|
||||
"""
|
||||
from app.services.local_diffusion import get_all_model_states
|
||||
|
||||
async def event_gen():
|
||||
try:
|
||||
while True:
|
||||
states = get_all_model_states()
|
||||
yield f"data: {json.dumps(states)}\n\n"
|
||||
await asyncio.sleep(0.2)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
return StreamingResponse(
|
||||
event_gen(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"X-Accel-Buffering": "no",
|
||||
"Connection": "keep-alive",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/generate/txt2img")
|
||||
async def txt2img(req: Txt2ImgRequest):
|
||||
"""Text-to-image via configured remote provider."""
|
||||
provider = _require_remote("txt2img")
|
||||
try:
|
||||
params = {
|
||||
"negative_prompt": req.negative_prompt or "",
|
||||
"steps": req.steps,
|
||||
"cfg_scale": req.cfg_scale,
|
||||
"seed": req.seed or 0,
|
||||
}
|
||||
if req.model:
|
||||
params["model"] = req.model
|
||||
result = await provider.txt2img(req.prompt, req.width, req.height, params)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/generate/img2img")
|
||||
async def img2img(req: Img2ImgRequest):
|
||||
"""Image-to-image via configured remote provider."""
|
||||
provider = _require_remote("img2img")
|
||||
try:
|
||||
params = {
|
||||
"negative_prompt": req.negative_prompt or "",
|
||||
"steps": req.steps,
|
||||
"cfg_scale": req.cfg_scale,
|
||||
}
|
||||
if req.model:
|
||||
params["model"] = req.model
|
||||
result = await provider.img2img(_decode(req.image), req.prompt, req.strength, params)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/generate/outpaint")
|
||||
async def outpaint(req: OutpaintRequest):
|
||||
"""Expand canvas in given direction via remote provider."""
|
||||
provider = _require_remote("outpaint")
|
||||
if req.direction not in ("left", "right", "top", "bottom"):
|
||||
raise HTTPException(status_code=400, detail="direction must be left/right/top/bottom")
|
||||
try:
|
||||
result = await provider.outpaint(_decode(req.image), req.direction, req.size, req.prompt or "")
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# ─── Config / capabilities ────────────────────────────────────────────────────
|
||||
|
||||
class ConfigUpdateRequest(BaseModel):
|
||||
ai_provider: Optional[str] = None
|
||||
# Per-operation overrides (blank = use default)
|
||||
ai_provider_inpaint: Optional[str] = None
|
||||
ai_provider_txt2img: Optional[str] = None
|
||||
ai_provider_img2img: Optional[str] = None
|
||||
ai_provider_outpaint: Optional[str] = None
|
||||
# Credentials / URLs
|
||||
openai_api_key: Optional[str] = None
|
||||
openai_model: Optional[str] = None
|
||||
invokeai_url: Optional[str] = None
|
||||
invokeai_default_model: Optional[str] = None
|
||||
comfyui_url: Optional[str] = None
|
||||
comfyui_default_model: Optional[str] = None
|
||||
replicate_api_key: Optional[str] = None
|
||||
stability_api_key: Optional[str] = None
|
||||
|
||||
|
||||
@router.post("/config")
|
||||
async def update_config(req: ConfigUpdateRequest):
|
||||
"""
|
||||
Apply runtime provider settings (no restart needed).
|
||||
Values are applied to the live settings object in-process.
|
||||
They do NOT persist across restarts — set them in .env for permanence.
|
||||
"""
|
||||
from app.config import settings
|
||||
|
||||
_str_fields = [
|
||||
"ai_provider", "ai_provider_inpaint", "ai_provider_txt2img",
|
||||
"ai_provider_img2img", "ai_provider_outpaint",
|
||||
"openai_api_key", "openai_model",
|
||||
"invokeai_url", "invokeai_default_model",
|
||||
"comfyui_url", "comfyui_default_model",
|
||||
"replicate_api_key", "stability_api_key",
|
||||
]
|
||||
for field in _str_fields:
|
||||
val = getattr(req, field, None)
|
||||
if val is not None:
|
||||
setattr(settings, field, val)
|
||||
|
||||
return {
|
||||
"status": "ok",
|
||||
"ai_provider": settings.ai_provider,
|
||||
"overrides": {
|
||||
"inpaint": settings.ai_provider_inpaint or None,
|
||||
"txt2img": settings.ai_provider_txt2img or None,
|
||||
"img2img": settings.ai_provider_img2img or None,
|
||||
"outpaint": settings.ai_provider_outpaint or None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
async def _check_provider(operation: str) -> dict:
|
||||
"""Health-check the provider for a specific operation."""
|
||||
from app.services.remote_provider import get_remote_provider
|
||||
try:
|
||||
p = get_remote_provider(operation)
|
||||
if p is None:
|
||||
return {"provider": None, "healthy": False}
|
||||
healthy = await asyncio.wait_for(p.health(), timeout=5.0)
|
||||
return {"provider": p.__class__.__name__.replace("Provider", "").lower(), "healthy": healthy}
|
||||
except Exception:
|
||||
return {"provider": None, "healthy": False}
|
||||
|
||||
|
||||
@router.get("/config")
|
||||
async def get_config():
|
||||
"""
|
||||
Return capability flags so the frontend can show/hide tools.
|
||||
Includes per-operation provider assignments and health status.
|
||||
"""
|
||||
from app.config import settings
|
||||
|
||||
# Run health checks for each operation concurrently
|
||||
ops = ["inpaint", "txt2img", "img2img", "outpaint"]
|
||||
results = await asyncio.gather(*[_check_provider(op) for op in ops])
|
||||
op_status = dict(zip(ops, results))
|
||||
|
||||
# Default provider for display (used when no per-op override)
|
||||
default_name = (settings.ai_provider or "").lower() or None
|
||||
|
||||
from app.services.gpu_detect import get_cached_gpu_info
|
||||
gpu_info = get_cached_gpu_info()
|
||||
|
||||
return {
|
||||
"local": {
|
||||
"lama": lama_available(),
|
||||
"rembg": rembg_available(),
|
||||
"opencv": True,
|
||||
"gpu_detected": gpu_available(),
|
||||
"gpu_backend": gpu_info.backend,
|
||||
"gpu_device": gpu_info.device_name,
|
||||
"gpu_vram_total": gpu_info.vram_total_gb,
|
||||
"gpu_vram_free": gpu_info.vram_free_gb,
|
||||
"gpu_cc": gpu_info.compute_capability,
|
||||
"gpu_fp16": gpu_info.fp16,
|
||||
"gpu_bf16": gpu_info.bf16,
|
||||
"gpu_fp8": gpu_info.fp8,
|
||||
"gpu_tensor_cores": gpu_info.tensor_cores,
|
||||
"gpu_tier": gpu_info.tier,
|
||||
"gpu_eff_vram": gpu_info.effective_vram_gb,
|
||||
"local_gpu_available": gpu_info.backend in ("cuda", "mps"),
|
||||
"local_gpu_capabilities": gpu_info.capabilities,
|
||||
"local_gpu_warnings": gpu_info.warnings,
|
||||
},
|
||||
"remote": {
|
||||
"default_provider": default_name,
|
||||
# Legacy field kept for backwards compat with badge/capabilities checks
|
||||
"provider": default_name,
|
||||
"healthy": any(v["healthy"] for v in op_status.values()),
|
||||
"operations": op_status,
|
||||
"overrides": {
|
||||
"inpaint": settings.ai_provider_inpaint or None,
|
||||
"txt2img": settings.ai_provider_txt2img or None,
|
||||
"img2img": settings.ai_provider_img2img or None,
|
||||
"outpaint": settings.ai_provider_outpaint or None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
# ─── Selection image operations ─────────────────────────────────────────────
|
||||
|
||||
class ScaleSelectionRequest(BaseModel):
|
||||
image: str # base64 full canvas
|
||||
mask: str # base64 selection mask (white = object)
|
||||
scale_pct: float = 103.0 # 103 = 3% bigger, 95 = 5% smaller
|
||||
|
||||
|
||||
class AiEditRegionRequest(BaseModel):
|
||||
image: str
|
||||
mask: str
|
||||
instruction: str
|
||||
negative_prompt: str = ""
|
||||
steps: int = 30
|
||||
cfg_scale: float = 7.5
|
||||
|
||||
|
||||
class PasteIntoSelectionRequest(BaseModel):
|
||||
image: str # base64 target canvas
|
||||
mask: str # base64 selection mask
|
||||
paste_image: str # base64 image to paste
|
||||
|
||||
|
||||
@router.post("/image/scale-selection")
|
||||
async def scale_selection(req: ScaleSelectionRequest):
|
||||
"""
|
||||
Scale the object selected by mask by scale_pct%, AI-fill the exposed gap.
|
||||
Works purely with local tools (LaMa/OpenCV) — no remote provider needed.
|
||||
"""
|
||||
try:
|
||||
import numpy as np
|
||||
from PIL import Image, ImageFilter
|
||||
except ImportError:
|
||||
raise HTTPException(status_code=500, detail="PIL/numpy not available")
|
||||
|
||||
img = Image.open(BytesIO(_decode(req.image))).convert("RGB")
|
||||
mask = Image.open(BytesIO(_decode(req.mask))).convert("L")
|
||||
if img.size != mask.size:
|
||||
mask = mask.resize(img.size, Image.LANCZOS)
|
||||
|
||||
mask_arr = np.array(mask)
|
||||
ys, xs = np.where(mask_arr > 128)
|
||||
if len(xs) == 0:
|
||||
raise HTTPException(status_code=400, detail="Empty mask — nothing to scale")
|
||||
|
||||
minx, maxx = int(xs.min()), int(xs.max())
|
||||
miny, maxy = int(ys.min()), int(ys.max())
|
||||
cx, cy = (minx + maxx) / 2.0, (miny + maxy) / 2.0
|
||||
obj_w, obj_h = maxx - minx + 1, maxy - miny + 1
|
||||
|
||||
scale = req.scale_pct / 100.0
|
||||
new_w = max(1, round(obj_w * scale))
|
||||
new_h = max(1, round(obj_h * scale))
|
||||
|
||||
# Extract masked object crop (RGBA with mask as alpha)
|
||||
img_rgba = img.convert("RGBA")
|
||||
obj_crop = img_rgba.crop((minx, miny, maxx + 1, maxy + 1))
|
||||
mask_crop = mask.crop((minx, miny, maxx + 1, maxy + 1))
|
||||
r, g, b, _ = obj_crop.split()
|
||||
obj_masked = Image.merge("RGBA", (r, g, b, mask_crop))
|
||||
scaled_obj = obj_masked.resize((new_w, new_h), Image.LANCZOS)
|
||||
|
||||
# AI-fill the original mask area (gap) with LaMa/OpenCV
|
||||
gap_mask = mask.filter(ImageFilter.MaxFilter(9)) # expand ~4px for clean seam
|
||||
gap_bytes = BytesIO()
|
||||
img.save(gap_bytes, format="PNG")
|
||||
gap_mask_bytes = BytesIO()
|
||||
gap_mask.save(gap_mask_bytes, format="PNG")
|
||||
|
||||
try:
|
||||
if lama_available():
|
||||
filled_bytes = await asyncio.get_event_loop().run_in_executor(
|
||||
None, lama_inpaint, gap_bytes.getvalue(), gap_mask_bytes.getvalue()
|
||||
)
|
||||
else:
|
||||
filled_bytes = await asyncio.get_event_loop().run_in_executor(
|
||||
None, opencv_inpaint, gap_bytes.getvalue(), gap_mask_bytes.getvalue()
|
||||
)
|
||||
filled = Image.open(BytesIO(filled_bytes)).convert("RGBA")
|
||||
except Exception as exc:
|
||||
print(f"[scale-selection] fill fallback: {exc}")
|
||||
filled = img.convert("RGBA")
|
||||
|
||||
# Paste scaled object centered on original centroid
|
||||
px = round(cx - new_w / 2)
|
||||
py = round(cy - new_h / 2)
|
||||
result = filled.copy()
|
||||
result.paste(scaled_obj, (px, py), scaled_obj.split()[3])
|
||||
|
||||
out = BytesIO()
|
||||
result.convert("RGB").save(out, format="PNG")
|
||||
return {"result": _encode(out.getvalue())}
|
||||
|
||||
|
||||
@router.post("/image/ai-edit-region")
|
||||
async def ai_edit_region(req: AiEditRegionRequest):
|
||||
"""
|
||||
AI-edit the selected region using the configured inpaint provider.
|
||||
Works with local_gpu, InvokeAI, ComfyUI, or OpenAI.
|
||||
"""
|
||||
provider = _require_remote("inpaint")
|
||||
try:
|
||||
result_bytes = await provider.inpaint(
|
||||
_decode(req.image),
|
||||
_decode(req.mask),
|
||||
req.instruction,
|
||||
{"negative_prompt": req.negative_prompt, "steps": req.steps, "cfg_scale": req.cfg_scale},
|
||||
)
|
||||
except Exception as exc:
|
||||
import traceback; traceback.print_exc()
|
||||
msg = str(exc)
|
||||
if "Errno -3" in msg or "Name or service not known" in msg or "ConnectError" in msg:
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail=(
|
||||
"AI model files not yet downloaded — container DNS appears to be blocked. "
|
||||
"Fix: sudo iptables -I DOCKER-USER -p udp --dport 53 -j ACCEPT on the host, "
|
||||
"or pre-download the model: pip install huggingface-hub && "
|
||||
"huggingface-cli download diffusers/stable-diffusion-xl-1.0-inpainting-0.1 "
|
||||
"--cache-dir ./data/hf_cache"
|
||||
)
|
||||
)
|
||||
raise HTTPException(status_code=500, detail=msg)
|
||||
return {"result": _encode(result_bytes)}
|
||||
|
||||
|
||||
@router.post("/image/paste-into-selection")
|
||||
async def paste_into_selection(req: PasteIntoSelectionRequest):
|
||||
"""
|
||||
Scale a clipboard image to the selection bounding box, mask it to the
|
||||
selection shape, and composite it over the original canvas.
|
||||
"""
|
||||
try:
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
except ImportError:
|
||||
raise HTTPException(status_code=500, detail="PIL/numpy not available")
|
||||
|
||||
img = Image.open(BytesIO(_decode(req.image))).convert("RGBA")
|
||||
mask = Image.open(BytesIO(_decode(req.mask))).convert("L")
|
||||
paste_img = Image.open(BytesIO(_decode(req.paste_image))).convert("RGBA")
|
||||
|
||||
if img.size != mask.size:
|
||||
mask = mask.resize(img.size, Image.LANCZOS)
|
||||
|
||||
mask_arr = np.array(mask)
|
||||
ys, xs = np.where(mask_arr > 128)
|
||||
if len(xs) == 0:
|
||||
raise HTTPException(status_code=400, detail="Empty mask")
|
||||
|
||||
minx, maxx = int(xs.min()), int(xs.max())
|
||||
miny, maxy = int(ys.min()), int(ys.max())
|
||||
target_w = maxx - minx + 1
|
||||
target_h = maxy - miny + 1
|
||||
|
||||
# Scale clipboard image to fit the selection bounding box
|
||||
paste_scaled = paste_img.resize((target_w, target_h), Image.LANCZOS)
|
||||
|
||||
# Clip paste to selection shape using mask
|
||||
mask_crop = mask.crop((minx, miny, maxx + 1, maxy + 1))
|
||||
r, g, b, a = paste_scaled.split()
|
||||
mask_np = np.array(mask_crop)
|
||||
alpha_np = np.array(a)
|
||||
combined = (alpha_np.astype(np.uint16) * mask_np.astype(np.uint16) // 255).astype(np.uint8)
|
||||
paste_final = Image.merge("RGBA", (r, g, b, Image.fromarray(combined)))
|
||||
|
||||
result = img.copy()
|
||||
result.paste(paste_final, (minx, miny), paste_final.split()[3])
|
||||
|
||||
out = BytesIO()
|
||||
result.convert("RGB").save(out, format="PNG")
|
||||
return {"result": _encode(out.getvalue())}
|
||||
|
||||
|
||||
# ─── SAM (Segment Anything) ──────────────────────────────────────────────────
|
||||
|
||||
class SegmentPointRequest(BaseModel):
|
||||
image: str # base64 PNG/JPEG
|
||||
points: list[list[int]] # [[x, y], ...] original image coords
|
||||
labels: list[int] # 1=include, 0=exclude — same length as points
|
||||
|
||||
|
||||
@router.post("/segment/point")
|
||||
async def segment_point(req: SegmentPointRequest):
|
||||
"""
|
||||
Run SAM point-prompt segmentation.
|
||||
Returns a binary mask PNG (white = selected area).
|
||||
Auto-downloads the SAM ViT-B model (~375 MB) on first call.
|
||||
"""
|
||||
if not req.points:
|
||||
raise HTTPException(status_code=400, detail="At least one point required.")
|
||||
if len(req.points) != len(req.labels):
|
||||
raise HTTPException(status_code=400, detail="points and labels must have the same length.")
|
||||
|
||||
try:
|
||||
image_bytes = base64.b64decode(req.image)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Could not decode image: {e}")
|
||||
|
||||
from app.services.sam_service import predict_points, get_install_status
|
||||
try:
|
||||
mask_bytes = await predict_points(
|
||||
image_bytes,
|
||||
[tuple(p) for p in req.points],
|
||||
req.labels,
|
||||
)
|
||||
return {
|
||||
"mask": base64.b64encode(mask_bytes).decode(),
|
||||
"sam_install": get_install_status(),
|
||||
}
|
||||
except RuntimeError as e:
|
||||
raise HTTPException(status_code=503, detail=str(e))
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/segment/install-status")
|
||||
def segment_install_status():
|
||||
"""Poll SAM model download progress."""
|
||||
from app.services.sam_service import get_install_status, sam_model_available
|
||||
status = get_install_status()
|
||||
status["model_ready"] = sam_model_available()
|
||||
return status
|
||||
|
||||
|
||||
@router.post("/segment/install")
|
||||
async def segment_install():
|
||||
"""Trigger SAM model download explicitly (also auto-triggered on first /segment/point call)."""
|
||||
from app.services.sam_service import ensure_sam_installed, get_install_status
|
||||
asyncio.create_task(ensure_sam_installed())
|
||||
return get_install_status()
|
||||
|
||||
|
||||
# ─── Enhance ─────────────────────────────────────────────────────────────────
|
||||
|
||||
import io as _io
|
||||
import numpy as _np
|
||||
import cv2 as _cv2
|
||||
from PIL import Image as _Image
|
||||
|
||||
class EnhanceRequest(BaseModel):
|
||||
image: str # base64
|
||||
strength: float = 1.0
|
||||
|
||||
|
||||
def _enhance_image(image_bytes: bytes, strength: float) -> bytes:
|
||||
"""
|
||||
Apply a chain of non-AI image enhancements, each blended with `strength` (0–1).
|
||||
|
||||
Steps:
|
||||
1. Auto white balance (gray-world)
|
||||
2. CLAHE on L channel of LAB colorspace
|
||||
3. Auto saturation boost in HSV (×1.15, clamped)
|
||||
4. Mild unsharp mask (gaussian sigma=1.0, delta weight=0.3)
|
||||
"""
|
||||
strength = max(0.0, min(1.0, float(strength)))
|
||||
|
||||
# Decode to RGB numpy array
|
||||
pil = _Image.open(_io.BytesIO(image_bytes)).convert("RGB")
|
||||
orig = _np.array(pil, dtype=_np.float32) # H×W×3, float [0,255]
|
||||
|
||||
img = orig.copy()
|
||||
|
||||
# ── Step 1: Auto white balance (gray-world) ──────────────────────────────
|
||||
mean_r = img[:, :, 0].mean()
|
||||
mean_g = img[:, :, 1].mean()
|
||||
mean_b = img[:, :, 2].mean()
|
||||
overall_mean = (mean_r + mean_g + mean_b) / 3.0
|
||||
|
||||
def _scale(channel, channel_mean):
|
||||
if channel_mean == 0:
|
||||
return channel
|
||||
return channel * (overall_mean / channel_mean)
|
||||
|
||||
wb = img.copy()
|
||||
wb[:, :, 0] = _np.clip(_scale(img[:, :, 0], mean_r), 0, 255)
|
||||
wb[:, :, 1] = _np.clip(_scale(img[:, :, 1], mean_g), 0, 255)
|
||||
wb[:, :, 2] = _np.clip(_scale(img[:, :, 2], mean_b), 0, 255)
|
||||
|
||||
img = (orig + strength * (wb - orig)).clip(0, 255)
|
||||
|
||||
# ── Step 2: CLAHE on L channel (LAB) ────────────────────────────────────
|
||||
img_u8 = img.astype(_np.uint8)
|
||||
lab = _cv2.cvtColor(img_u8, _cv2.COLOR_RGB2LAB)
|
||||
clahe = _cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||||
l_orig = lab[:, :, 0].copy()
|
||||
lab[:, :, 0] = clahe.apply(l_orig)
|
||||
# Blend L channel back using strength
|
||||
lab_blended = lab.copy()
|
||||
lab_blended[:, :, 0] = (l_orig + strength * (lab[:, :, 0].astype(_np.float32) - l_orig.astype(_np.float32))).clip(0, 255).astype(_np.uint8)
|
||||
img = _cv2.cvtColor(lab_blended, _cv2.COLOR_LAB2RGB).astype(_np.float32)
|
||||
|
||||
# ── Step 3: Auto saturation boost (HSV, ×1.15) ──────────────────────────
|
||||
img_u8 = img.astype(_np.uint8)
|
||||
hsv = _cv2.cvtColor(img_u8, _cv2.COLOR_RGB2HSV).astype(_np.float32)
|
||||
s_orig = hsv[:, :, 1].copy()
|
||||
s_boosted = _np.clip(s_orig * 1.15, 0, 255)
|
||||
hsv[:, :, 1] = s_orig + strength * (s_boosted - s_orig)
|
||||
hsv = hsv.clip(0, 255).astype(_np.uint8)
|
||||
img = _cv2.cvtColor(hsv, _cv2.COLOR_HSV2RGB).astype(_np.float32)
|
||||
|
||||
# ── Step 4: Mild unsharp mask (sigma=1.0, delta weight=0.3) ─────────────
|
||||
img_u8 = img.astype(_np.uint8)
|
||||
blurred = _cv2.GaussianBlur(img_u8, (0, 0), sigmaX=1.0)
|
||||
sharpness_delta = img_u8.astype(_np.float32) - blurred.astype(_np.float32)
|
||||
sharpened = img_u8.astype(_np.float32) + 0.3 * sharpness_delta * strength
|
||||
img = sharpened.clip(0, 255)
|
||||
|
||||
# Encode result as PNG
|
||||
result_pil = _Image.fromarray(img.astype(_np.uint8), mode="RGB")
|
||||
buf = _io.BytesIO()
|
||||
result_pil.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
@router.post("/enhance")
|
||||
async def enhance(req: EnhanceRequest):
|
||||
"""
|
||||
Non-AI image enhancement: auto white balance, CLAHE, saturation boost,
|
||||
and unsharp mask. Each step is blended proportionally to `strength` (0–1).
|
||||
"""
|
||||
try:
|
||||
image_bytes = _decode(req.image)
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, _enhance_image, image_bytes, req.strength
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# ─── Subject replace ─────────────────────────────────────────────────────────
|
||||
|
||||
class ExtractSubjectRequest(BaseModel):
|
||||
image: str # base64
|
||||
|
||||
|
||||
class ReplaceSubjectRequest(BaseModel):
|
||||
background_image: str # base64 — image whose background we keep
|
||||
subject_image: str # base64 — image whose subject we extract
|
||||
mask: Optional[str] = None # base64 — white = where the subject should land
|
||||
match_colors: bool = True # blend subject color stats toward background
|
||||
|
||||
|
||||
def _extract_subject_bytes(image_bytes: bytes) -> bytes:
|
||||
"""Remove background from image using rembg; return RGBA PNG bytes."""
|
||||
if rembg_available():
|
||||
return remove_background_rembg(image_bytes)
|
||||
raise RuntimeError(
|
||||
"rembg is not installed. Run: pip install rembg (or add it to requirements.txt)"
|
||||
)
|
||||
|
||||
|
||||
def _color_transfer_lab(subj_rgba: "Image", bg_rgb: "Image", blend: float = 0.45) -> "Image":
|
||||
"""
|
||||
Partial LAB color transfer: nudge subject color statistics 'blend' fraction
|
||||
toward the background's statistics so it looks like it belongs in the scene.
|
||||
"""
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
src_arr = np.array(subj_rgba.convert("RGB"), dtype=np.float32)
|
||||
tgt_arr = np.array(bg_rgb.convert("RGB"), dtype=np.float32)
|
||||
|
||||
alpha = np.array(subj_rgba.split()[3])
|
||||
subject_mask = alpha > 10
|
||||
|
||||
if not subject_mask.any():
|
||||
return subj_rgba
|
||||
|
||||
src_lab = cv2.cvtColor(src_arr.astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
|
||||
tgt_lab = cv2.cvtColor(tgt_arr.astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
|
||||
|
||||
for ch in range(3):
|
||||
src_ch = src_lab[:, :, ch]
|
||||
src_pixels = src_ch[subject_mask]
|
||||
tgt_pixels = tgt_lab[:, :, ch].flatten()
|
||||
|
||||
src_mean, src_std = float(src_pixels.mean()), float(src_pixels.std()) + 1e-6
|
||||
tgt_mean, tgt_std = float(tgt_pixels.mean()), float(tgt_pixels.std()) + 1e-6
|
||||
|
||||
adjusted_std = src_std + blend * (tgt_std - src_std)
|
||||
adjusted = (src_ch - src_mean) * (adjusted_std / src_std) + src_mean + blend * (tgt_mean - src_mean)
|
||||
src_lab[:, :, ch] = np.clip(adjusted, 0, 255)
|
||||
|
||||
result_rgb = cv2.cvtColor(src_lab.astype(np.uint8), cv2.COLOR_LAB2RGB)
|
||||
r, g, b = result_rgb[:, :, 0], result_rgb[:, :, 1], result_rgb[:, :, 2]
|
||||
return Image.merge("RGBA", [
|
||||
Image.fromarray(r), Image.fromarray(g),
|
||||
Image.fromarray(b), Image.fromarray(alpha),
|
||||
])
|
||||
|
||||
|
||||
def _do_replace_subject(
|
||||
bg_bytes: bytes,
|
||||
subj_bytes: bytes,
|
||||
mask_bytes: Optional[bytes],
|
||||
match_colors: bool,
|
||||
) -> bytes:
|
||||
"""Core compositing: extract subject → scale → color-match → paste onto background."""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
bg_img = Image.open(BytesIO(bg_bytes)).convert("RGBA")
|
||||
|
||||
subj_rgba = Image.open(BytesIO(_extract_subject_bytes(subj_bytes))).convert("RGBA")
|
||||
|
||||
# Determine target placement bounding box from mask or full canvas
|
||||
if mask_bytes:
|
||||
mask_img = Image.open(BytesIO(mask_bytes)).convert("L")
|
||||
if mask_img.size != bg_img.size:
|
||||
mask_img = mask_img.resize(bg_img.size, Image.LANCZOS)
|
||||
mask_arr = np.array(mask_img)
|
||||
ys, xs = np.where(mask_arr > 128)
|
||||
else:
|
||||
mask_img = None
|
||||
mask_arr = None
|
||||
ys, xs = np.array([]), np.array([])
|
||||
|
||||
if len(xs) > 0:
|
||||
minx, maxx = int(xs.min()), int(xs.max())
|
||||
miny, maxy = int(ys.min()), int(ys.max())
|
||||
else:
|
||||
minx, miny = 0, 0
|
||||
maxx, maxy = bg_img.width - 1, bg_img.height - 1
|
||||
|
||||
target_w = maxx - minx + 1
|
||||
target_h = maxy - miny + 1
|
||||
|
||||
# Scale subject to fit target area, preserving aspect ratio
|
||||
sw, sh = subj_rgba.size
|
||||
scale = min(target_w / sw, target_h / sh)
|
||||
new_w = max(1, round(sw * scale))
|
||||
new_h = max(1, round(sh * scale))
|
||||
subj_scaled = subj_rgba.resize((new_w, new_h), Image.LANCZOS)
|
||||
|
||||
# Optional color transfer to blend lighting/tone
|
||||
if match_colors:
|
||||
subj_scaled = _color_transfer_lab(subj_scaled, bg_img.convert("RGB"))
|
||||
|
||||
# Center in target area
|
||||
px = minx + (target_w - new_w) // 2
|
||||
py = miny + (target_h - new_h) // 2
|
||||
|
||||
result = bg_img.copy()
|
||||
|
||||
if mask_img is not None and len(xs) > 0:
|
||||
# Build a full-canvas RGBA layer for the subject
|
||||
subj_canvas = Image.new("RGBA", bg_img.size, (0, 0, 0, 0))
|
||||
subj_canvas.paste(subj_scaled, (px, py), subj_scaled.split()[3])
|
||||
# Clip subject's alpha to the selection mask
|
||||
sc_arr = np.array(subj_canvas)
|
||||
sc_arr[:, :, 3] = np.minimum(sc_arr[:, :, 3], mask_arr).astype(np.uint8)
|
||||
subj_canvas = Image.fromarray(sc_arr)
|
||||
result.paste(subj_canvas, (0, 0), subj_canvas.split()[3])
|
||||
else:
|
||||
result.paste(subj_scaled, (px, py), subj_scaled.split()[3])
|
||||
|
||||
out = BytesIO()
|
||||
result.convert("RGB").save(out, format="PNG")
|
||||
return out.getvalue()
|
||||
|
||||
|
||||
@router.post("/image/extract-subject")
|
||||
async def extract_subject(req: ExtractSubjectRequest):
|
||||
"""
|
||||
Remove background from an image and return the subject with transparency (RGBA PNG).
|
||||
Uses rembg (AI-powered) when available.
|
||||
"""
|
||||
try:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, _extract_subject_bytes, _decode(req.image)
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/image/replace-subject")
|
||||
async def replace_subject(req: ReplaceSubjectRequest):
|
||||
"""
|
||||
Extract the primary subject from `subject_image` (via rembg background removal),
|
||||
scale it to fit the `mask` selection on `background_image`, apply optional LAB
|
||||
color transfer for lighting consistency, and composite the result.
|
||||
|
||||
Returns the composited image as base64 PNG.
|
||||
"""
|
||||
try:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None,
|
||||
_do_replace_subject,
|
||||
_decode(req.background_image),
|
||||
_decode(req.subject_image),
|
||||
_decode(req.mask) if req.mask else None,
|
||||
req.match_colors,
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# ─── Extract colors ───────────────────────────────────────────────────────────
|
||||
|
||||
class ExtractColorsRequest(BaseModel):
|
||||
image: str # base64
|
||||
count: int = 6
|
||||
|
||||
|
||||
def _extract_colors(image_bytes: bytes, count: int) -> list[str]:
|
||||
"""
|
||||
Resize image to 150×150, k-means cluster pixels into `count` groups
|
||||
using pure numpy (no sklearn dependency), return hex strings by frequency.
|
||||
"""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from io import BytesIO
|
||||
|
||||
count = max(1, min(count, 32))
|
||||
|
||||
pil = Image.open(BytesIO(image_bytes)).convert("RGB").resize((150, 150))
|
||||
pixels = np.array(pil, dtype=np.float32).reshape(-1, 3) # (22500, 3)
|
||||
n = len(pixels)
|
||||
|
||||
# Initialise centers with k-means++ seeding
|
||||
rng = np.random.default_rng(42)
|
||||
centers = [pixels[rng.integers(n)]]
|
||||
for _ in range(count - 1):
|
||||
dists = np.min([np.sum((pixels - c) ** 2, axis=1) for c in centers], axis=0)
|
||||
probs = dists / dists.sum()
|
||||
centers.append(pixels[rng.choice(n, p=probs)])
|
||||
centers = np.array(centers)
|
||||
|
||||
labels = np.zeros(n, dtype=np.int32)
|
||||
for _ in range(20): # max 20 iterations
|
||||
# Assign each pixel to nearest center
|
||||
dists = np.sum((pixels[:, None] - centers[None]) ** 2, axis=2) # (n, k)
|
||||
new_labels = np.argmin(dists, axis=1)
|
||||
if np.all(new_labels == labels):
|
||||
break
|
||||
labels = new_labels
|
||||
# Recompute centers
|
||||
for k in range(count):
|
||||
mask = labels == k
|
||||
if mask.any():
|
||||
centers[k] = pixels[mask].mean(axis=0)
|
||||
|
||||
counts = np.bincount(labels, minlength=count)
|
||||
order = np.argsort(-counts)
|
||||
|
||||
return [
|
||||
"#{:02x}{:02x}{:02x}".format(*centers[i].astype(int).clip(0, 255))
|
||||
for i in order
|
||||
]
|
||||
|
||||
|
||||
@router.post("/extract-colors")
|
||||
async def extract_colors(req: ExtractColorsRequest):
|
||||
"""
|
||||
Extract dominant colors from an image using k-means clustering.
|
||||
Returns hex color strings sorted by frequency (most dominant first).
|
||||
"""
|
||||
try:
|
||||
image_bytes = _decode(req.image)
|
||||
colors = await asyncio.get_event_loop().run_in_executor(
|
||||
None, _extract_colors, image_bytes, req.count
|
||||
)
|
||||
return {"colors": colors}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
@@ -0,0 +1,171 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks
|
||||
from sqlalchemy.orm import Session
|
||||
import json
|
||||
|
||||
from app.database import get_db
|
||||
from app.models.project import Project
|
||||
from app.models.edit import Edit
|
||||
from app.schemas import EditRequest, EditResponse, StatusResponse
|
||||
from app.services.edit_service import EditService
|
||||
from app.config import settings
|
||||
|
||||
router = APIRouter(prefix="/edits", tags=["edits"])
|
||||
|
||||
|
||||
async def process_edit_background(
|
||||
edit_id: int,
|
||||
project_id: int,
|
||||
request: EditRequest,
|
||||
db: Session
|
||||
):
|
||||
"""Background task to process edit"""
|
||||
edit_service = EditService()
|
||||
|
||||
try:
|
||||
# Process the edit
|
||||
result_path = await edit_service.process_edit(
|
||||
project_id=project_id,
|
||||
edit_id=edit_id,
|
||||
prompt=request.prompt,
|
||||
mode=request.mode,
|
||||
selection_type=request.selection_type,
|
||||
bbox=request.bbox,
|
||||
feather_px=request.feather_px,
|
||||
selection_data=request.selection_data
|
||||
)
|
||||
|
||||
# Update edit status
|
||||
edit = db.query(Edit).filter(Edit.id == edit_id).first()
|
||||
if edit:
|
||||
edit.status = "completed"
|
||||
db.commit()
|
||||
|
||||
except Exception as e:
|
||||
# Update edit with error
|
||||
edit = db.query(Edit).filter(Edit.id == edit_id).first()
|
||||
if edit:
|
||||
edit.status = "failed"
|
||||
edit.error_message = str(e)
|
||||
db.commit()
|
||||
|
||||
|
||||
@router.post("/projects/{project_id}/fix", response_model=EditResponse)
|
||||
async def create_edit(
|
||||
project_id: int,
|
||||
request: EditRequest,
|
||||
background_tasks: BackgroundTasks,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Create a new edit request (Fix button)
|
||||
|
||||
This endpoint accepts the selection data and prompt,
|
||||
then processes the edit in the background.
|
||||
"""
|
||||
# Verify project exists
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Validate mode
|
||||
if request.mode not in ["A", "B"]:
|
||||
raise HTTPException(status_code=400, detail="Mode must be 'A' or 'B'")
|
||||
|
||||
# Validate selection type
|
||||
if request.selection_type not in ["rectangle", "ellipse", "lasso"]:
|
||||
raise HTTPException(status_code=400, detail="Invalid selection type")
|
||||
|
||||
# Create edit record
|
||||
edit = Edit(
|
||||
project_id=project_id,
|
||||
mode=request.mode,
|
||||
prompt=request.prompt,
|
||||
selection_type=request.selection_type,
|
||||
bbox_json=json.dumps(request.bbox),
|
||||
feather_px=request.feather_px,
|
||||
ai_provider=settings.ai_provider,
|
||||
status="pending"
|
||||
)
|
||||
db.add(edit)
|
||||
db.commit()
|
||||
db.refresh(edit)
|
||||
|
||||
# Process edit in background
|
||||
background_tasks.add_task(
|
||||
process_edit_background,
|
||||
edit.id,
|
||||
project_id,
|
||||
request,
|
||||
db
|
||||
)
|
||||
|
||||
return edit
|
||||
|
||||
|
||||
@router.get("/{edit_id}", response_model=EditResponse)
|
||||
def get_edit(
|
||||
edit_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get edit details and status"""
|
||||
edit = db.query(Edit).filter(Edit.id == edit_id).first()
|
||||
if not edit:
|
||||
raise HTTPException(status_code=404, detail="Edit not found")
|
||||
return edit
|
||||
|
||||
|
||||
@router.post("/projects/{project_id}/revert/{edit_id}", response_model=StatusResponse)
|
||||
def revert_to_edit(
|
||||
project_id: int,
|
||||
edit_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Revert project to a specific edit"""
|
||||
# Verify project exists
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Verify edit exists and belongs to project
|
||||
edit = db.query(Edit).filter(
|
||||
Edit.id == edit_id,
|
||||
Edit.project_id == project_id
|
||||
).first()
|
||||
if not edit:
|
||||
raise HTTPException(status_code=404, detail="Edit not found")
|
||||
|
||||
# Revert
|
||||
edit_service = EditService()
|
||||
try:
|
||||
result_path = edit_service.revert_to_edit(project_id, edit_id)
|
||||
return StatusResponse(
|
||||
status="success",
|
||||
message=f"Reverted to edit {edit_id}",
|
||||
data={"image_url": f"/projects/{project_id}/current"}
|
||||
)
|
||||
except FileNotFoundError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/projects/{project_id}/reset", response_model=StatusResponse)
|
||||
def reset_to_original(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Reset project to original image"""
|
||||
# Verify project exists
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Reset
|
||||
edit_service = EditService()
|
||||
try:
|
||||
result_path = edit_service.reset_to_original(project_id)
|
||||
return StatusResponse(
|
||||
status="success",
|
||||
message="Reset to original image",
|
||||
data={"image_url": f"/projects/{project_id}/current"}
|
||||
)
|
||||
except FileNotFoundError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
@@ -0,0 +1,176 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, Form
|
||||
from sqlalchemy.orm import Session
|
||||
from typing import Optional
|
||||
from PIL import Image
|
||||
from io import BytesIO
|
||||
import os
|
||||
|
||||
from app.database import get_db
|
||||
from app.models.project import Project
|
||||
from app.schemas import TextToImageRequest, TextToImageResponse
|
||||
from app.services.ai_provider import get_ai_provider
|
||||
from app.services.edit_service import EditService
|
||||
from app.config import settings
|
||||
|
||||
router = APIRouter(prefix="/generate", tags=["generate"])
|
||||
|
||||
|
||||
@router.post("/text-to-image", response_model=TextToImageResponse)
|
||||
async def text_to_image(
|
||||
prompt: str = Form(...),
|
||||
width: int = Form(1024),
|
||||
height: int = Form(1024),
|
||||
negative_prompt: Optional[str] = Form(None),
|
||||
ai_provider: Optional[str] = Form(None),
|
||||
ai_model: Optional[str] = Form(None),
|
||||
create_project: bool = Form(True),
|
||||
project_name: Optional[str] = Form(None),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Generate an image from text prompt
|
||||
|
||||
Args:
|
||||
prompt: Text description of desired image
|
||||
width: Image width (default 1024)
|
||||
height: Image height (default 1024)
|
||||
negative_prompt: What to avoid in generation
|
||||
ai_provider: Override default AI provider
|
||||
ai_model: Specific model to use
|
||||
create_project: Whether to create a new project with the result
|
||||
project_name: Name for the new project (if create_project=True)
|
||||
|
||||
Returns:
|
||||
Generated image info and optionally project details
|
||||
"""
|
||||
|
||||
# Validate dimensions
|
||||
if width < 256 or width > 2048 or height < 256 or height > 2048:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Width and height must be between 256 and 2048"
|
||||
)
|
||||
|
||||
# Get AI provider
|
||||
provider = get_ai_provider(ai_provider, ai_model)
|
||||
|
||||
try:
|
||||
# Generate image
|
||||
image_bytes = await provider.text_to_image(
|
||||
prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
model=ai_model,
|
||||
negative_prompt=negative_prompt
|
||||
)
|
||||
|
||||
project_id = None
|
||||
image_url = None
|
||||
|
||||
if create_project:
|
||||
# Create a new project
|
||||
project = Project(
|
||||
name=project_name or f"Generated: {prompt[:50]}",
|
||||
user_id=None # TODO: Add authentication
|
||||
)
|
||||
db.add(project)
|
||||
db.commit()
|
||||
db.refresh(project)
|
||||
|
||||
project_id = project.id
|
||||
|
||||
# Save image as both original and current
|
||||
edit_service = EditService()
|
||||
edit_service.ensure_project_dir(project_id)
|
||||
|
||||
original_path = edit_service.get_original_image_path(project_id)
|
||||
current_path = edit_service.get_current_image_path(project_id)
|
||||
|
||||
# Save image
|
||||
img = Image.open(BytesIO(image_bytes))
|
||||
img.save(original_path, 'PNG')
|
||||
img.save(current_path, 'PNG')
|
||||
|
||||
image_url = f"/projects/{project_id}/current"
|
||||
|
||||
return TextToImageResponse(
|
||||
status="success",
|
||||
prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
project_id=project_id,
|
||||
image_url=image_url,
|
||||
ai_provider=ai_provider or settings.ai_provider,
|
||||
ai_model=ai_model
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/layer/text-to-image", response_model=TextToImageResponse)
|
||||
async def text_to_image_layer(
|
||||
project_id: int = Form(...),
|
||||
prompt: str = Form(...),
|
||||
width: int = Form(512),
|
||||
height: int = Form(512),
|
||||
x: int = Form(0),
|
||||
y: int = Form(0),
|
||||
negative_prompt: Optional[str] = Form(None),
|
||||
ai_provider: Optional[str] = Form(None),
|
||||
ai_model: Optional[str] = Form(None),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Generate an image as a new layer in an existing project
|
||||
|
||||
This generates a smaller image that can be placed as a layer
|
||||
on top of the current project canvas.
|
||||
"""
|
||||
|
||||
# Verify project exists
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Get AI provider
|
||||
provider = get_ai_provider(ai_provider, ai_model)
|
||||
|
||||
try:
|
||||
# Generate image
|
||||
image_bytes = await provider.text_to_image(
|
||||
prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
model=ai_model,
|
||||
negative_prompt=negative_prompt
|
||||
)
|
||||
|
||||
# Save as temporary layer file
|
||||
edit_service = EditService()
|
||||
layers_dir = edit_service.get_project_dir(project_id) / "layers"
|
||||
layers_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Generate unique layer filename
|
||||
import time
|
||||
layer_filename = f"generated_{int(time.time())}.png"
|
||||
layer_path = layers_dir / layer_filename
|
||||
|
||||
# Save layer image
|
||||
with open(layer_path, 'wb') as f:
|
||||
f.write(image_bytes)
|
||||
|
||||
return TextToImageResponse(
|
||||
status="success",
|
||||
prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
project_id=project_id,
|
||||
image_url=f"/projects/{project_id}/layers/{layer_filename}",
|
||||
layer_position={"x": x, "y": y},
|
||||
ai_provider=ai_provider or settings.ai_provider,
|
||||
ai_model=ai_model
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
@@ -0,0 +1,96 @@
|
||||
"""
|
||||
GPU status and model management endpoints.
|
||||
All under /api/gpu prefix.
|
||||
"""
|
||||
from fastapi import APIRouter
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, List
|
||||
import asyncio
|
||||
|
||||
router = APIRouter(prefix="/api/gpu", tags=["gpu"])
|
||||
|
||||
|
||||
@router.get("/status")
|
||||
async def gpu_status():
|
||||
"""
|
||||
Full GPU capability report: hardware, feature flags, VRAM budget,
|
||||
and which model was selected for each operation.
|
||||
Frontend polls this to show GPU badge and tool availability.
|
||||
"""
|
||||
from app.services.gpu_detect import get_cached_gpu_info
|
||||
from app.services.local_diffusion import get_all_model_states
|
||||
|
||||
info = get_cached_gpu_info()
|
||||
|
||||
return {
|
||||
# Hardware
|
||||
"backend": info.backend,
|
||||
"device_name": info.device_name,
|
||||
"vram_total_gb": info.vram_total_gb,
|
||||
"vram_free_gb": info.vram_free_gb,
|
||||
"compute_capability": info.compute_capability,
|
||||
# Feature flags
|
||||
"fp16": info.fp16,
|
||||
"bf16": info.bf16,
|
||||
"fp8": info.fp8,
|
||||
"int8": info.int8,
|
||||
"tensor_cores": info.tensor_cores,
|
||||
"xformers": info.xformers,
|
||||
# Derived
|
||||
"effective_vram_gb": info.effective_vram_gb,
|
||||
"tier": info.tier,
|
||||
# Selected models per operation
|
||||
"recommended": {
|
||||
op: (
|
||||
{
|
||||
"model_id": spec.model_id,
|
||||
"family": spec.family,
|
||||
"memory_opt": spec.memory_opt,
|
||||
"native_res": spec.native_res,
|
||||
"vram_fp16_gb": spec.vram_fp16_gb,
|
||||
}
|
||||
if spec else None
|
||||
)
|
||||
for op, spec in info.recommended.items()
|
||||
},
|
||||
"pipeline_states": get_all_model_states(),
|
||||
"warnings": info.warnings,
|
||||
"capabilities": info.capabilities,
|
||||
}
|
||||
|
||||
|
||||
class PrefetchRequest(BaseModel):
|
||||
operations: Optional[List[str]] = None
|
||||
|
||||
|
||||
@router.post("/prefetch")
|
||||
async def prefetch_models(req: PrefetchRequest = PrefetchRequest()):
|
||||
"""
|
||||
Eagerly load pipelines into GPU memory for the requested operations.
|
||||
Returns immediately; poll /api/gpu/prefetch-status for progress.
|
||||
Default: inpaint, txt2img, img2img.
|
||||
"""
|
||||
ops = req.operations or ["inpaint", "txt2img", "img2img"]
|
||||
valid = {"inpaint", "txt2img", "img2img", "outpaint", "upscale"}
|
||||
ops = [op for op in ops if op in valid]
|
||||
|
||||
from app.services.local_diffusion import get_local_diffusion_provider
|
||||
provider = get_local_diffusion_provider()
|
||||
|
||||
async def _prefetch():
|
||||
for op in ops:
|
||||
try:
|
||||
await provider._get_pipeline(op)
|
||||
print(f"[gpu] Prefetch complete: {op}")
|
||||
except Exception as exc:
|
||||
print(f"[gpu] Prefetch failed for {op}: {exc}")
|
||||
|
||||
asyncio.create_task(_prefetch())
|
||||
return {"status": "prefetch_started", "operations": ops}
|
||||
|
||||
|
||||
@router.get("/prefetch-status")
|
||||
async def prefetch_status():
|
||||
"""Poll model download / load progress."""
|
||||
from app.services.local_diffusion import get_all_model_states
|
||||
return {"models": get_all_model_states()}
|
||||
@@ -0,0 +1,81 @@
|
||||
from fastapi import APIRouter, HTTPException, Depends
|
||||
from fastapi.responses import FileResponse
|
||||
from sqlalchemy.orm import Session
|
||||
from pathlib import Path
|
||||
|
||||
from app.database import get_db
|
||||
from app.models.project import Project
|
||||
from app.services.edit_service import EditService
|
||||
|
||||
router = APIRouter(prefix="/projects", tags=["images"])
|
||||
|
||||
|
||||
@router.get("/{project_id}/original")
|
||||
def get_original_image(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get the original uploaded image"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
edit_service = EditService()
|
||||
image_path = edit_service.get_original_image_path(project_id)
|
||||
|
||||
if not image_path.exists():
|
||||
raise HTTPException(status_code=404, detail="Original image not found")
|
||||
|
||||
return FileResponse(
|
||||
image_path,
|
||||
media_type="image/png",
|
||||
headers={"Cache-Control": "public, max-age=3600"}
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{project_id}/current")
|
||||
def get_current_image(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get the current edited image"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
edit_service = EditService()
|
||||
image_path = edit_service.get_current_image_path(project_id)
|
||||
|
||||
if not image_path.exists():
|
||||
raise HTTPException(status_code=404, detail="Current image not found")
|
||||
|
||||
return FileResponse(
|
||||
image_path,
|
||||
media_type="image/png",
|
||||
headers={"Cache-Control": "no-cache"}
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{project_id}/history/{edit_id}/result")
|
||||
def get_edit_result(
|
||||
project_id: int,
|
||||
edit_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get the result image from a specific edit"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
edit_service = EditService()
|
||||
edit_dir = edit_service.get_edit_dir(project_id, edit_id)
|
||||
result_path = edit_dir / "result.png"
|
||||
|
||||
if not result_path.exists():
|
||||
raise HTTPException(status_code=404, detail="Edit result not found")
|
||||
|
||||
return FileResponse(
|
||||
result_path,
|
||||
media_type="image/png",
|
||||
headers={"Cache-Control": "public, max-age=3600"}
|
||||
)
|
||||
@@ -0,0 +1,308 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form
|
||||
from fastapi.responses import FileResponse
|
||||
from sqlalchemy.orm import Session
|
||||
from typing import List, Optional
|
||||
import json
|
||||
|
||||
from app.database import get_db
|
||||
from app.models.patch import Patch
|
||||
from app.models.project import Project
|
||||
from app.models.edit import Edit
|
||||
from app.schemas import PatchCreate, PatchResponse, PatchApply, StatusResponse
|
||||
from app.services.patch_library import PatchLibraryService
|
||||
from app.config import settings
|
||||
|
||||
router = APIRouter(prefix="/patches", tags=["patches"])
|
||||
|
||||
|
||||
@router.post("/", response_model=PatchResponse)
|
||||
async def create_patch(
|
||||
name: str = Form(...),
|
||||
description: Optional[str] = Form(None),
|
||||
source_type: str = Form(...),
|
||||
category: Optional[str] = Form(None),
|
||||
tags: Optional[str] = Form(None),
|
||||
source_project_id: Optional[int] = Form(None),
|
||||
source_edit_id: Optional[int] = Form(None),
|
||||
bbox: Optional[str] = Form(None),
|
||||
file: Optional[UploadFile] = File(None),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Create a new patch in the library
|
||||
|
||||
Source types:
|
||||
- ai_generated: From an edit (requires source_edit_id)
|
||||
- manual_selection: Selected from current image (requires source_project_id and bbox)
|
||||
- imported: Uploaded file (requires file)
|
||||
"""
|
||||
|
||||
# Validate source_type
|
||||
if source_type not in ["ai_generated", "manual_selection", "imported"]:
|
||||
raise HTTPException(status_code=400, detail="Invalid source_type")
|
||||
|
||||
# Create patch record
|
||||
patch = Patch(
|
||||
name=name,
|
||||
description=description,
|
||||
source_type=source_type,
|
||||
source_project_id=source_project_id,
|
||||
source_edit_id=source_edit_id,
|
||||
tags=tags,
|
||||
category=category,
|
||||
file_path="", # Will be set after saving
|
||||
user_id=None # TODO: Add authentication
|
||||
)
|
||||
|
||||
db.add(patch)
|
||||
db.commit()
|
||||
db.refresh(patch)
|
||||
|
||||
# Save patch file based on source type
|
||||
patch_service = PatchLibraryService()
|
||||
|
||||
try:
|
||||
if source_type == "ai_generated":
|
||||
# Get edit directory and save AI-generated patch
|
||||
if not source_edit_id:
|
||||
raise HTTPException(status_code=400, detail="source_edit_id required for ai_generated")
|
||||
|
||||
edit = db.query(Edit).filter(Edit.id == source_edit_id).first()
|
||||
if not edit:
|
||||
raise HTTPException(status_code=404, detail="Edit not found")
|
||||
|
||||
from app.services.edit_service import EditService
|
||||
edit_service = EditService()
|
||||
edit_dir = edit_service.get_edit_dir(edit.project_id, edit.id)
|
||||
|
||||
file_path = patch_service.save_ai_generated_patch(patch.id, edit_dir)
|
||||
|
||||
# Get dimensions
|
||||
width, height = patch_service.get_patch_size(patch.id)
|
||||
patch.width = width
|
||||
patch.height = height
|
||||
|
||||
elif source_type == "manual_selection":
|
||||
# Save manually selected patch from current image
|
||||
if not source_project_id or not bbox:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="source_project_id and bbox required for manual_selection"
|
||||
)
|
||||
|
||||
project = db.query(Project).filter(Project.id == source_project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
bbox_dict = json.loads(bbox) if isinstance(bbox, str) else bbox
|
||||
file_path = patch_service.save_manual_patch(patch.id, source_project_id, bbox_dict)
|
||||
|
||||
patch.width = bbox_dict['width']
|
||||
patch.height = bbox_dict['height']
|
||||
|
||||
elif source_type == "imported":
|
||||
# Save uploaded file
|
||||
if not file:
|
||||
raise HTTPException(status_code=400, detail="file required for imported")
|
||||
|
||||
image_bytes = await file.read()
|
||||
file_path = patch_service.save_patch_from_bytes(patch.id, image_bytes)
|
||||
|
||||
# Get dimensions
|
||||
width, height = patch_service.get_patch_size(patch.id)
|
||||
patch.width = width
|
||||
patch.height = height
|
||||
|
||||
# Update patch with file path
|
||||
patch.file_path = file_path
|
||||
patch.thumbnail_path = str(patch_service.get_thumbnail_path(patch.id))
|
||||
db.commit()
|
||||
db.refresh(patch)
|
||||
|
||||
return patch
|
||||
|
||||
except Exception as e:
|
||||
# Cleanup on error
|
||||
patch_service.delete_patch(patch.id)
|
||||
db.delete(patch)
|
||||
db.commit()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/", response_model=List[PatchResponse])
|
||||
def list_patches(
|
||||
category: Optional[str] = None,
|
||||
tags: Optional[str] = None,
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""List patches in the library with optional filtering"""
|
||||
|
||||
query = db.query(Patch)
|
||||
|
||||
if category:
|
||||
query = query.filter(Patch.category == category)
|
||||
|
||||
if tags:
|
||||
# Simple tag search (could be improved with full-text search)
|
||||
query = query.filter(Patch.tags.like(f"%{tags}%"))
|
||||
|
||||
patches = query.offset(offset).limit(limit).all()
|
||||
return patches
|
||||
|
||||
|
||||
@router.get("/{patch_id}", response_model=PatchResponse)
|
||||
def get_patch(
|
||||
patch_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get patch details"""
|
||||
patch = db.query(Patch).filter(Patch.id == patch_id).first()
|
||||
if not patch:
|
||||
raise HTTPException(status_code=404, detail="Patch not found")
|
||||
return patch
|
||||
|
||||
|
||||
@router.get("/{patch_id}/image")
|
||||
def get_patch_image(
|
||||
patch_id: int,
|
||||
thumbnail: bool = False,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get patch image file"""
|
||||
patch = db.query(Patch).filter(Patch.id == patch_id).first()
|
||||
if not patch:
|
||||
raise HTTPException(status_code=404, detail="Patch not found")
|
||||
|
||||
patch_service = PatchLibraryService()
|
||||
|
||||
if thumbnail:
|
||||
file_path = patch_service.get_thumbnail_path(patch_id)
|
||||
else:
|
||||
file_path = patch_service.get_patch_path(patch_id)
|
||||
|
||||
if not file_path.exists():
|
||||
raise HTTPException(status_code=404, detail="Patch image not found")
|
||||
|
||||
return FileResponse(file_path, media_type="image/png")
|
||||
|
||||
|
||||
@router.post("/apply", response_model=StatusResponse)
|
||||
async def apply_patch(
|
||||
project_id: int = Form(...),
|
||||
patch_id: int = Form(...),
|
||||
bbox: str = Form(...),
|
||||
feather_px: int = Form(5),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Apply a saved patch to a project image
|
||||
|
||||
This creates a new edit in the project history.
|
||||
"""
|
||||
# Verify project exists
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Verify patch exists
|
||||
patch = db.query(Patch).filter(Patch.id == patch_id).first()
|
||||
if not patch:
|
||||
raise HTTPException(status_code=404, detail="Patch not found")
|
||||
|
||||
# Parse bbox
|
||||
bbox_dict = json.loads(bbox) if isinstance(bbox, str) else bbox
|
||||
|
||||
# Load current image
|
||||
from app.services.edit_service import EditService
|
||||
from PIL import Image
|
||||
|
||||
edit_service = EditService()
|
||||
current_image_path = edit_service.get_current_image_path(project_id)
|
||||
current_image = Image.open(current_image_path).convert('RGBA')
|
||||
|
||||
# Apply patch
|
||||
patch_service = PatchLibraryService()
|
||||
result_image = patch_service.apply_patch_to_image(
|
||||
patch_id,
|
||||
current_image,
|
||||
bbox_dict,
|
||||
feather_px
|
||||
)
|
||||
|
||||
# Save result as current image
|
||||
result_image.save(current_image_path)
|
||||
|
||||
# Create edit record
|
||||
edit = Edit(
|
||||
project_id=project_id,
|
||||
mode="patch_library",
|
||||
prompt=f"Applied saved patch: {patch.name}",
|
||||
selection_type="rectangle",
|
||||
bbox_json=json.dumps(bbox_dict),
|
||||
feather_px=feather_px,
|
||||
ai_provider="patch_library",
|
||||
status="completed"
|
||||
)
|
||||
db.add(edit)
|
||||
db.commit()
|
||||
|
||||
return StatusResponse(
|
||||
status="success",
|
||||
message=f"Applied patch '{patch.name}' to project",
|
||||
data={"edit_id": edit.id}
|
||||
)
|
||||
|
||||
|
||||
@router.delete("/{patch_id}", response_model=StatusResponse)
|
||||
def delete_patch(
|
||||
patch_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Delete a patch from the library"""
|
||||
patch = db.query(Patch).filter(Patch.id == patch_id).first()
|
||||
if not patch:
|
||||
raise HTTPException(status_code=404, detail="Patch not found")
|
||||
|
||||
# Delete files
|
||||
patch_service = PatchLibraryService()
|
||||
patch_service.delete_patch(patch_id)
|
||||
|
||||
# Delete record
|
||||
db.delete(patch)
|
||||
db.commit()
|
||||
|
||||
return StatusResponse(
|
||||
status="success",
|
||||
message=f"Deleted patch '{patch.name}'"
|
||||
)
|
||||
|
||||
|
||||
@router.put("/{patch_id}", response_model=PatchResponse)
|
||||
def update_patch(
|
||||
patch_id: int,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
category: Optional[str] = None,
|
||||
tags: Optional[str] = None,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Update patch metadata"""
|
||||
patch = db.query(Patch).filter(Patch.id == patch_id).first()
|
||||
if not patch:
|
||||
raise HTTPException(status_code=404, detail="Patch not found")
|
||||
|
||||
if name:
|
||||
patch.name = name
|
||||
if description is not None:
|
||||
patch.description = description
|
||||
if category:
|
||||
patch.category = category
|
||||
if tags is not None:
|
||||
patch.tags = tags
|
||||
|
||||
db.commit()
|
||||
db.refresh(patch)
|
||||
|
||||
return patch
|
||||
@@ -0,0 +1,487 @@
|
||||
"""
|
||||
Print / frame tools — frame fit and upscale.
|
||||
All endpoints under /api/print prefix.
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, Literal
|
||||
import base64
|
||||
import asyncio
|
||||
from io import BytesIO
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
|
||||
router = APIRouter(prefix="/api/print", tags=["print-tools"])
|
||||
|
||||
# ── Frame size catalogue (inches) ──────────────────────────────────────────
|
||||
FRAME_SIZES = {
|
||||
"4x6": (4, 6),
|
||||
"5x7": (5, 7),
|
||||
"8x10": (8, 10),
|
||||
"11x14": (11, 14),
|
||||
"16x20": (16, 20),
|
||||
"18x24": (18, 24),
|
||||
"20x24": (20, 24),
|
||||
"24x36": (24, 36),
|
||||
# Square
|
||||
"4x4": (4, 4),
|
||||
"8x8": (8, 8),
|
||||
"12x12": (12, 12),
|
||||
}
|
||||
|
||||
|
||||
def _encode(data: bytes) -> str:
|
||||
return base64.b64encode(data).decode()
|
||||
|
||||
|
||||
def _decode(b64: str) -> bytes:
|
||||
return base64.b64decode(b64)
|
||||
|
||||
|
||||
def _to_png(img: Image.Image) -> bytes:
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
# ── Request models ─────────────────────────────────────────────────────────
|
||||
|
||||
class FrameFitRequest(BaseModel):
|
||||
image: str # base64 PNG/JPEG
|
||||
frame: str # e.g. "8x10"
|
||||
orientation: Literal["auto", "portrait", "landscape"] = "auto"
|
||||
mode: Literal["crop", "extend", "smart"] = "smart"
|
||||
dpi: int = 300
|
||||
# For extend mode: prompt passed to outpaint
|
||||
prompt: Optional[str] = ""
|
||||
# Smart mode threshold: extend if gap fraction < this, else crop
|
||||
smart_threshold: float = 0.15
|
||||
|
||||
|
||||
class UpscaleRequest(BaseModel):
|
||||
image: str # base64
|
||||
scale: float = 2.0 # 1.5, 2, 3, 4
|
||||
# auto = pick best available; lanczos = always works; realesrgan_pytorch / realesrgan_ncnn = explicit
|
||||
method: str = "auto"
|
||||
|
||||
|
||||
class PrepareRequest(BaseModel):
|
||||
image: str # base64
|
||||
frame: str # e.g. "8x10"
|
||||
orientation: Literal["auto", "portrait", "landscape"] = "auto"
|
||||
target_dpi: int = 300
|
||||
upscale_method: str = "auto" # auto / realesrgan_pytorch / realesrgan_ncnn / lanczos
|
||||
mode: Literal["crop", "extend", "smart"] = "smart"
|
||||
prompt: Optional[str] = ""
|
||||
|
||||
|
||||
# ── Frame sizes endpoint ───────────────────────────────────────────────────
|
||||
|
||||
@router.get("/frame-sizes")
|
||||
def list_frame_sizes():
|
||||
"""Return the catalogue of supported frame sizes."""
|
||||
return {
|
||||
"sizes": list(FRAME_SIZES.keys()),
|
||||
"catalogue": {k: {"inches": v, "pixels_300dpi": (v[0]*300, v[1]*300)}
|
||||
for k, v in FRAME_SIZES.items()},
|
||||
}
|
||||
|
||||
|
||||
# ── Frame fit ──────────────────────────────────────────────────────────────
|
||||
|
||||
@router.post("/frame-fit")
|
||||
async def frame_fit(req: FrameFitRequest):
|
||||
"""
|
||||
Fit an image to a print frame size.
|
||||
|
||||
Modes:
|
||||
crop — center-crop to frame aspect ratio, then scale to print resolution.
|
||||
extend — scale to fill one dimension, outpaint the gap with AI.
|
||||
smart — extend if gap < smart_threshold of frame dimension, else crop.
|
||||
|
||||
Returns the fitted image plus a summary of what was done.
|
||||
"""
|
||||
if req.frame not in FRAME_SIZES:
|
||||
raise HTTPException(status_code=400,
|
||||
detail=f"Unknown frame '{req.frame}'. Valid: {list(FRAME_SIZES.keys())}")
|
||||
if not (72 <= req.dpi <= 600):
|
||||
raise HTTPException(status_code=400, detail="dpi must be 72–600")
|
||||
|
||||
try:
|
||||
image = Image.open(BytesIO(_decode(req.image))).convert("RGB")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Could not decode image: {e}")
|
||||
|
||||
fw, fh = FRAME_SIZES[req.frame] # frame inches (w, h in portrait)
|
||||
|
||||
# Resolve orientation
|
||||
img_w, img_h = image.size
|
||||
img_landscape = img_w >= img_h
|
||||
frame_landscape = fw >= fh
|
||||
|
||||
if req.orientation == "landscape":
|
||||
fw, fh = max(fw, fh), min(fw, fh)
|
||||
elif req.orientation == "portrait":
|
||||
fw, fh = min(fw, fh), max(fw, fh)
|
||||
else: # auto — match image orientation
|
||||
if img_landscape and not frame_landscape:
|
||||
fw, fh = fh, fw # rotate frame to landscape
|
||||
elif not img_landscape and frame_landscape:
|
||||
fw, fh = fh, fw # rotate frame to portrait
|
||||
|
||||
target_w = fw * req.dpi
|
||||
target_h = fh * req.dpi
|
||||
target_ratio = target_w / target_h
|
||||
img_ratio = img_w / img_h
|
||||
|
||||
# Determine actual mode
|
||||
mode = req.mode
|
||||
if mode == "smart":
|
||||
# Scale image to fill the frame — compute gap fraction
|
||||
if img_ratio > target_ratio:
|
||||
# Image wider → fits on height, gap on width
|
||||
scaled_h = target_h
|
||||
scaled_w = round(target_h * img_ratio)
|
||||
gap_frac = (scaled_w - target_w) / target_w # positive = overflow (crop)
|
||||
else:
|
||||
scaled_w = target_w
|
||||
scaled_h = round(target_w / img_ratio)
|
||||
gap_frac = (scaled_h - target_h) / target_h
|
||||
|
||||
# gap_frac > 0 means we'd need to crop; < 0 means we'd need to extend
|
||||
if gap_frac < 0:
|
||||
# Need to extend — use extend if gap is small enough
|
||||
mode = "extend" if abs(gap_frac) <= req.smart_threshold else "crop"
|
||||
else:
|
||||
mode = "crop"
|
||||
|
||||
if mode == "crop":
|
||||
result, summary = _crop_fit(image, target_w, target_h)
|
||||
else: # extend
|
||||
result, summary = await _extend_fit(image, target_w, target_h, req.prompt or "")
|
||||
|
||||
return {
|
||||
"result": _encode(_to_png(result)),
|
||||
"mode_used": mode,
|
||||
"frame": req.frame,
|
||||
"orientation": "landscape" if fw > fh else "portrait",
|
||||
"output_pixels": {"width": result.width, "height": result.height},
|
||||
"output_inches": {"width": fw, "height": fh},
|
||||
"dpi": req.dpi,
|
||||
"summary": summary,
|
||||
}
|
||||
|
||||
|
||||
def _crop_fit(image: Image.Image, target_w: int, target_h: int):
|
||||
"""Center-crop image to target aspect ratio, then Lanczos scale to target size."""
|
||||
img_w, img_h = image.size
|
||||
target_ratio = target_w / target_h
|
||||
img_ratio = img_w / img_h
|
||||
|
||||
if img_ratio > target_ratio:
|
||||
# Wider than target — crop sides
|
||||
new_w = round(img_h * target_ratio)
|
||||
x0 = (img_w - new_w) // 2
|
||||
cropped = image.crop((x0, 0, x0 + new_w, img_h))
|
||||
else:
|
||||
# Taller than target — crop top/bottom
|
||||
new_h = round(img_w / target_ratio)
|
||||
y0 = (img_h - new_h) // 2
|
||||
cropped = image.crop((0, y0, img_w, y0 + new_h))
|
||||
|
||||
result = cropped.resize((target_w, target_h), Image.Resampling.LANCZOS)
|
||||
summary = (
|
||||
f"Cropped from {img_w}×{img_h} to {cropped.width}×{cropped.height}, "
|
||||
f"scaled to {target_w}×{target_h}"
|
||||
)
|
||||
return result, summary
|
||||
|
||||
|
||||
async def _extend_fit(image: Image.Image, target_w: int, target_h: int, prompt: str):
|
||||
"""
|
||||
Scale image to fill one dimension exactly, then outpaint the gap with AI.
|
||||
Falls back to content-aware mirror fill if no remote provider configured.
|
||||
"""
|
||||
from app.services.remote_provider import get_remote_provider
|
||||
|
||||
img_w, img_h = image.size
|
||||
target_ratio = target_w / target_h
|
||||
img_ratio = img_w / img_h
|
||||
|
||||
if img_ratio > target_ratio:
|
||||
# Image wider — scale to target width, extend height
|
||||
scale = target_w / img_w
|
||||
scaled_w = target_w
|
||||
scaled_h = round(img_h * scale)
|
||||
gap_dir = "height"
|
||||
gap_top = (target_h - scaled_h) // 2
|
||||
gap_bottom = target_h - scaled_h - gap_top
|
||||
else:
|
||||
# Image taller — scale to target height, extend width
|
||||
scale = target_h / img_h
|
||||
scaled_h = target_h
|
||||
scaled_w = round(img_w * scale)
|
||||
gap_dir = "width"
|
||||
gap_left = (target_w - scaled_w) // 2
|
||||
gap_right = target_w - scaled_w - gap_left
|
||||
|
||||
scaled = image.resize((scaled_w, scaled_h), Image.Resampling.LANCZOS)
|
||||
|
||||
# Place scaled image on canvas
|
||||
canvas = Image.new("RGB", (target_w, target_h), (128, 128, 128))
|
||||
if gap_dir == "height":
|
||||
canvas.paste(scaled, (0, gap_top))
|
||||
# Build mask: top and bottom strips are white (to inpaint)
|
||||
mask = Image.new("L", (target_w, target_h), 0)
|
||||
if gap_top > 0:
|
||||
mask.paste(Image.new("L", (target_w, gap_top), 255), (0, 0))
|
||||
if gap_bottom > 0:
|
||||
mask.paste(Image.new("L", (target_w, gap_bottom), 255), (0, target_h - gap_bottom))
|
||||
else:
|
||||
canvas.paste(scaled, (gap_left, 0))
|
||||
mask = Image.new("L", (target_w, target_h), 0)
|
||||
if gap_left > 0:
|
||||
mask.paste(Image.new("L", (gap_left, target_h), 255), (0, 0))
|
||||
if gap_right > 0:
|
||||
mask.paste(Image.new("L", (gap_right, target_h), 255), (target_w - gap_right, 0))
|
||||
|
||||
# Try AI inpaint
|
||||
provider = get_remote_provider("inpaint")
|
||||
if provider:
|
||||
try:
|
||||
canvas_bytes = _to_png(canvas)
|
||||
mask_bytes = _to_png(mask)
|
||||
fill_prompt = prompt or "seamlessly continue the image, natural extension"
|
||||
result_bytes = await provider.inpaint(canvas_bytes, mask_bytes, fill_prompt, {})
|
||||
result = Image.open(BytesIO(result_bytes)).convert("RGB")
|
||||
summary = (
|
||||
f"Scaled {img_w}×{img_h} → {scaled_w}×{scaled_h}, "
|
||||
f"AI-extended {gap_dir} to {target_w}×{target_h}"
|
||||
)
|
||||
return result, summary
|
||||
except Exception as e:
|
||||
print(f"AI extend failed, using mirror fill: {e}")
|
||||
|
||||
# Fallback: mirror-fill the gap (looks decent for backgrounds/landscapes)
|
||||
result = _mirror_fill(canvas, mask, scaled, gap_dir,
|
||||
gap_top if gap_dir == "height" else gap_left,
|
||||
gap_bottom if gap_dir == "height" else gap_right,
|
||||
target_w, target_h)
|
||||
summary = (
|
||||
f"Scaled {img_w}×{img_h} → {scaled_w}×{scaled_h}, "
|
||||
f"mirror-filled {gap_dir} to {target_w}×{target_h} (no AI provider)"
|
||||
)
|
||||
return result, summary
|
||||
|
||||
|
||||
def _mirror_fill(canvas, mask, scaled, gap_dir, gap_a, gap_b, target_w, target_h):
|
||||
"""Fill gaps by reflecting the nearest edge strip."""
|
||||
result = canvas.copy()
|
||||
if gap_dir == "height":
|
||||
if gap_a > 0:
|
||||
strip = scaled.crop((0, 0, scaled.width, min(gap_a * 2, scaled.height)))
|
||||
strip = strip.transpose(Image.Transpose.FLIP_TOP_BOTTOM)
|
||||
strip = strip.resize((target_w, gap_a), Image.Resampling.LANCZOS)
|
||||
result.paste(strip, (0, 0))
|
||||
if gap_b > 0:
|
||||
strip = scaled.crop((0, max(0, scaled.height - gap_b * 2), scaled.width, scaled.height))
|
||||
strip = strip.transpose(Image.Transpose.FLIP_TOP_BOTTOM)
|
||||
strip = strip.resize((target_w, gap_b), Image.Resampling.LANCZOS)
|
||||
result.paste(strip, (0, target_h - gap_b))
|
||||
else:
|
||||
if gap_a > 0:
|
||||
strip = scaled.crop((0, 0, min(gap_a * 2, scaled.width), scaled.height))
|
||||
strip = strip.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
|
||||
strip = strip.resize((gap_a, target_h), Image.Resampling.LANCZOS)
|
||||
result.paste(strip, (0, 0))
|
||||
if gap_b > 0:
|
||||
strip = scaled.crop((max(0, scaled.width - gap_b * 2), 0, scaled.width, scaled.height))
|
||||
strip = strip.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
|
||||
strip = strip.resize((gap_b, target_h), Image.Resampling.LANCZOS)
|
||||
result.paste(strip, (target_w - gap_b, 0))
|
||||
return result
|
||||
|
||||
|
||||
# ── Upscale ────────────────────────────────────────────────────────────────
|
||||
|
||||
@router.post("/upscale/refresh-caps")
|
||||
def upscale_refresh_caps():
|
||||
"""Bust the capability cache (call after installing Real-ESRGAN without restarting)."""
|
||||
from app.services.upscale import invalidate_caps_cache, probe_upscale_capabilities
|
||||
invalidate_caps_cache()
|
||||
return probe_upscale_capabilities()
|
||||
|
||||
|
||||
@router.get("/upscale/available")
|
||||
async def upscale_available():
|
||||
"""
|
||||
Return capability probe: which upscale methods are available,
|
||||
which device will be used, and which method is recommended.
|
||||
If no AI upscaler is found, triggers background NCNN auto-install.
|
||||
Frontend uses this to populate the method selector.
|
||||
"""
|
||||
from app.services.upscale import probe_upscale_capabilities, ensure_ncnn_installed, get_install_status
|
||||
caps = probe_upscale_capabilities()
|
||||
# Auto-install NCNN if no AI upscaler is available yet
|
||||
if not caps["realesrgan_pytorch"] and not caps["realesrgan_ncnn"]:
|
||||
asyncio.create_task(ensure_ncnn_installed())
|
||||
caps["ncnn_install_status"] = get_install_status()
|
||||
return caps
|
||||
|
||||
|
||||
@router.get("/upscale/install-status")
|
||||
def upscale_install_status():
|
||||
"""Poll for Real-ESRGAN NCNN auto-install progress."""
|
||||
from app.services.upscale import get_install_status, probe_upscale_capabilities, _find_ncnn_binary
|
||||
status = get_install_status()
|
||||
# If install just finished, refresh caps
|
||||
if status["state"] == "done":
|
||||
from app.services.upscale import invalidate_caps_cache
|
||||
invalidate_caps_cache()
|
||||
caps = probe_upscale_capabilities()
|
||||
status["ncnn_available"] = caps["realesrgan_ncnn"]
|
||||
else:
|
||||
status["ncnn_available"] = False
|
||||
return status
|
||||
|
||||
|
||||
@router.post("/prepare")
|
||||
async def prepare_for_print(req: PrepareRequest):
|
||||
"""
|
||||
One-shot Prepare for Print: AI upscale to reach target DPI, then fit to frame.
|
||||
|
||||
Steps:
|
||||
1. Resolve target pixel dimensions (frame × target_dpi, orientation-adjusted)
|
||||
2. Calculate needed upscale factor so the image meets the target resolution
|
||||
3. Run Real-ESRGAN if scale > 1.05 (else skip — already large enough)
|
||||
4. Run frame-fit (crop / extend / smart) to exact target dimensions
|
||||
5. Return the print-ready image and a quality report
|
||||
"""
|
||||
if req.frame not in FRAME_SIZES:
|
||||
raise HTTPException(status_code=400,
|
||||
detail=f"Unknown frame '{req.frame}'. Valid: {list(FRAME_SIZES.keys())}")
|
||||
if not (72 <= req.target_dpi <= 600):
|
||||
raise HTTPException(status_code=400, detail="target_dpi must be 72–600")
|
||||
|
||||
try:
|
||||
image = Image.open(BytesIO(_decode(req.image))).convert("RGB")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Could not decode image: {e}")
|
||||
|
||||
fw, fh = FRAME_SIZES[req.frame]
|
||||
img_w, img_h = image.size
|
||||
|
||||
# Resolve orientation (same logic as frame_fit)
|
||||
img_landscape = img_w >= img_h
|
||||
frame_landscape = fw >= fh
|
||||
if req.orientation == "landscape":
|
||||
fw, fh = max(fw, fh), min(fw, fh)
|
||||
elif req.orientation == "portrait":
|
||||
fw, fh = min(fw, fh), max(fw, fh)
|
||||
else:
|
||||
if img_landscape and not frame_landscape:
|
||||
fw, fh = fh, fw
|
||||
elif not img_landscape and frame_landscape:
|
||||
fw, fh = fh, fw
|
||||
|
||||
target_w = fw * req.target_dpi
|
||||
target_h = fh * req.target_dpi
|
||||
|
||||
# Scale factor needed so the shorter dimension fills the frame
|
||||
scale_w = target_w / img_w
|
||||
scale_h = target_h / img_h
|
||||
needed_scale = min(scale_w, scale_h) # fill-to-fit (extend) baseline
|
||||
# For crop mode we need max; use the larger to be safe and let frame-fit crop
|
||||
needed_scale_crop = max(scale_w, scale_h)
|
||||
|
||||
# Use the smaller (extend) scale as the upscale target; frame-fit handles the rest
|
||||
upscale_factor = max(1.0, needed_scale)
|
||||
upscale_applied = False
|
||||
method_used = "none"
|
||||
|
||||
upscaled = image
|
||||
if upscale_factor > 1.05:
|
||||
# Cap per-pass at 4× (Real-ESRGAN works best at 2–4×)
|
||||
remaining = upscale_factor
|
||||
while remaining > 1.05:
|
||||
pass_scale = min(remaining, 4.0)
|
||||
# Round to one decimal to keep scale in 1.1–8.0 range accepted by upscale service
|
||||
pass_scale = round(pass_scale, 1)
|
||||
if pass_scale < 1.1:
|
||||
break
|
||||
from app.services.upscale import upscale_image
|
||||
result_bytes, method_used = await upscale_image(upscaled, pass_scale, req.upscale_method)
|
||||
upscaled = Image.open(BytesIO(result_bytes)).convert("RGB")
|
||||
remaining /= pass_scale
|
||||
upscale_applied = True
|
||||
|
||||
# Encode upscaled image and run frame-fit
|
||||
upscaled_b64 = _encode(_to_png(upscaled))
|
||||
|
||||
fit_req = FrameFitRequest(
|
||||
image=upscaled_b64,
|
||||
frame=req.frame,
|
||||
orientation=req.orientation,
|
||||
mode=req.mode,
|
||||
dpi=req.target_dpi,
|
||||
prompt=req.prompt or "",
|
||||
)
|
||||
# Re-use the existing frame_fit logic inline
|
||||
fit_response = await frame_fit(fit_req)
|
||||
|
||||
return {
|
||||
"result": fit_response["result"],
|
||||
"frame": req.frame,
|
||||
"orientation": fit_response["orientation"],
|
||||
"output_pixels": fit_response["output_pixels"],
|
||||
"output_inches": fit_response["output_inches"],
|
||||
"dpi": req.target_dpi,
|
||||
"mode_used": fit_response["mode_used"],
|
||||
"upscale_applied": upscale_applied,
|
||||
"upscale_factor": round(upscale_factor, 2),
|
||||
"upscale_method": method_used,
|
||||
"summary": fit_response["summary"],
|
||||
}
|
||||
|
||||
|
||||
@router.post("/upscale")
|
||||
async def upscale(req: UpscaleRequest):
|
||||
"""
|
||||
Upscale image. method values:
|
||||
auto — pick best available (recommended)
|
||||
realesrgan_pytorch — Real-ESRGAN via PyTorch (CUDA/MPS/CPU)
|
||||
realesrgan_ncnn — Real-ESRGAN NCNN Vulkan binary
|
||||
lanczos — always available, instant
|
||||
Any AI method falls back to the next best if unavailable.
|
||||
"""
|
||||
if not (1.1 <= req.scale <= 8.0):
|
||||
raise HTTPException(status_code=400, detail="scale must be 1.1–8.0")
|
||||
|
||||
valid_methods = {"auto", "realesrgan_pytorch", "realesrgan_ncnn", "lanczos"}
|
||||
if req.method not in valid_methods:
|
||||
raise HTTPException(status_code=400,
|
||||
detail=f"method must be one of {sorted(valid_methods)}")
|
||||
|
||||
try:
|
||||
image = Image.open(BytesIO(_decode(req.image))).convert("RGB")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Could not decode image: {e}")
|
||||
|
||||
orig_w, orig_h = image.size
|
||||
|
||||
try:
|
||||
from app.services.upscale import upscale_image
|
||||
result_bytes, method_used = await upscale_image(image, req.scale, req.method)
|
||||
result = Image.open(BytesIO(result_bytes))
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
return {
|
||||
"result": _encode(result_bytes),
|
||||
"method": method_used,
|
||||
"original": {"width": orig_w, "height": orig_h},
|
||||
"output": {"width": result.width, "height": result.height},
|
||||
"scale": req.scale,
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, status
|
||||
from sqlalchemy.orm import Session
|
||||
from typing import List
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from PIL import Image
|
||||
|
||||
from app.database import get_db
|
||||
from app.models.project import Project
|
||||
from app.models.edit import Edit
|
||||
from app.schemas import ProjectCreate, ProjectResponse, EditResponse, UploadResponse
|
||||
from app.services.edit_service import EditService
|
||||
from app.config import settings
|
||||
|
||||
router = APIRouter(prefix="/projects", tags=["projects"])
|
||||
|
||||
|
||||
@router.post("/", response_model=ProjectResponse)
|
||||
def create_project(
|
||||
project: ProjectCreate,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Create a new project"""
|
||||
# For MVP, we'll use a default user_id of 1
|
||||
# In production, this would come from authentication
|
||||
user_id = 1
|
||||
|
||||
db_project = Project(
|
||||
user_id=user_id,
|
||||
name=project.name
|
||||
)
|
||||
db.add(db_project)
|
||||
db.commit()
|
||||
db.refresh(db_project)
|
||||
|
||||
# Create project directory
|
||||
edit_service = EditService()
|
||||
edit_service.ensure_project_dir(db_project.id)
|
||||
|
||||
return db_project
|
||||
|
||||
|
||||
@router.get("/", response_model=List[ProjectResponse])
|
||||
def list_projects(
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""List all projects"""
|
||||
projects = db.query(Project).offset(skip).limit(limit).all()
|
||||
return projects
|
||||
|
||||
|
||||
@router.get("/{project_id}", response_model=ProjectResponse)
|
||||
def get_project(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get a specific project"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
return project
|
||||
|
||||
|
||||
@router.delete("/{project_id}")
|
||||
def delete_project(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Delete a project"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Delete project directory
|
||||
edit_service = EditService()
|
||||
project_dir = edit_service.get_project_dir(project_id)
|
||||
if project_dir.exists():
|
||||
shutil.rmtree(project_dir)
|
||||
|
||||
db.delete(project)
|
||||
db.commit()
|
||||
|
||||
return {"status": "success", "message": f"Project {project_id} deleted"}
|
||||
|
||||
|
||||
@router.post("/{project_id}/upload", response_model=UploadResponse)
|
||||
async def upload_image(
|
||||
project_id: int,
|
||||
file: UploadFile = File(...),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Upload an image to a project"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Validate file type
|
||||
if not file.content_type.startswith('image/'):
|
||||
raise HTTPException(status_code=400, detail="File must be an image")
|
||||
|
||||
# Create project directory
|
||||
edit_service = EditService()
|
||||
edit_service.ensure_project_dir(project_id)
|
||||
|
||||
# Save original and current images
|
||||
original_path = edit_service.get_original_image_path(project_id)
|
||||
current_path = edit_service.get_current_image_path(project_id)
|
||||
|
||||
# Read and validate image
|
||||
contents = await file.read()
|
||||
try:
|
||||
image = Image.open(BytesIO(contents))
|
||||
image = image.convert('RGBA')
|
||||
|
||||
# Save images
|
||||
image.save(original_path, 'PNG')
|
||||
image.save(current_path, 'PNG')
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Invalid image file: {str(e)}")
|
||||
|
||||
return UploadResponse(
|
||||
project_id=project_id,
|
||||
original_url=f"/projects/{project_id}/original",
|
||||
current_url=f"/projects/{project_id}/current"
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{project_id}/edits", response_model=List[EditResponse])
|
||||
def list_edits(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""List all edits for a project"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
edits = db.query(Edit).filter(Edit.project_id == project_id).order_by(Edit.created_at.desc()).all()
|
||||
return edits
|
||||
|
||||
|
||||
from io import BytesIO
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user