Backend: - requirements.txt: add simple-lama-inpainting, rembg[gpu]; upgrade opencv to 4.10+ - app/config.py: add InvokeAI (url, model) and ComfyUI (url, model) settings; OPENAI_MODEL - app/services/local_inpaint.py: LaMa, OpenCV, rembg wrappers (auto GPU/CPU) - app/services/remote_provider.py: abstract RemoteAIProvider + OpenAI, InvokeAI, ComfyUI drivers - app/routers/ai_tools.py: new /api/* endpoints — /erase, /inpaint/lama, /inpaint/fast, /background/remove, /inpaint/remote, /generate/txt2img, /generate/img2img, /generate/outpaint, GET /config (capability flags) - app/main.py: register ai_tools router Frontend: - services/api.js: add erase(), textToImage(), imageToImage(), remoteInpaint(), getConfig() - api/capabilities.js: lazy-fetch /api/config singleton; hasRemote() helper - tools/ai_lama_erase.js: brush-paint mask → LaMa erase → apply to layer - tools/ai_smart_inpaint.js: brush mask + dialog (Fast/Quality mode + prompt) → inpaint - core/components/provider-badge.js: shows active provider + health in toolbar - config.js: register ai_lama_erase and ai_smart_inpaint tools - main.js: mount provider badge on load - .env.example: document InvokeAI, ComfyUI, OpenAI provider settings https://claude.ai/code/session_01B58MaJCU1R6KwBDJCp8AfN
83 lines
2.3 KiB
Python
83 lines
2.3 KiB
Python
"""
|
|
Local inpainting operations — LaMa, OpenCV, and background removal.
|
|
All operations use GPU automatically if PyTorch detects one, CPU otherwise.
|
|
"""
|
|
|
|
from io import BytesIO
|
|
from PIL import Image
|
|
import numpy as np
|
|
import cv2
|
|
|
|
# Lazy-loaded LaMa model (downloaded on first use, ~100MB)
|
|
_lama = None
|
|
|
|
|
|
def get_lama():
|
|
global _lama
|
|
if _lama is None:
|
|
from simple_lama_inpainting import SimpleLama
|
|
_lama = SimpleLama()
|
|
return _lama
|
|
|
|
|
|
def lama_available() -> bool:
|
|
try:
|
|
import simple_lama_inpainting # noqa: F401
|
|
return True
|
|
except ImportError:
|
|
return False
|
|
|
|
|
|
def lama_inpaint(image_bytes: bytes, mask_bytes: bytes) -> bytes:
|
|
"""LaMa structural inpainting — best for object removal and large fills."""
|
|
lama = get_lama()
|
|
image = Image.open(BytesIO(image_bytes)).convert("RGB")
|
|
mask = Image.open(BytesIO(mask_bytes)).convert("L")
|
|
if mask.size != image.size:
|
|
mask = mask.resize(image.size, Image.Resampling.LANCZOS)
|
|
result = lama(image, mask)
|
|
buf = BytesIO()
|
|
result.save(buf, format="PNG")
|
|
return buf.getvalue()
|
|
|
|
|
|
def opencv_inpaint(image_bytes: bytes, mask_bytes: bytes, method: str = "telea") -> bytes:
|
|
"""OpenCV fast structural inpainting — CPU only, milliseconds."""
|
|
image = Image.open(BytesIO(image_bytes)).convert("RGB")
|
|
mask = Image.open(BytesIO(mask_bytes)).convert("L")
|
|
if mask.size != image.size:
|
|
mask = mask.resize(image.size, Image.Resampling.LANCZOS)
|
|
|
|
img_np = np.array(image)
|
|
mask_np = np.array(mask)
|
|
_, mask_bin = cv2.threshold(mask_np, 127, 255, cv2.THRESH_BINARY)
|
|
|
|
flags = cv2.INPAINT_TELEA if method == "telea" else cv2.INPAINT_NS
|
|
result = cv2.inpaint(img_np, mask_bin, inpaintRadius=3, flags=flags)
|
|
|
|
buf = BytesIO()
|
|
Image.fromarray(result).save(buf, format="PNG")
|
|
return buf.getvalue()
|
|
|
|
|
|
def remove_background_rembg(image_bytes: bytes) -> bytes:
|
|
"""Background removal using rembg."""
|
|
from rembg import remove
|
|
return remove(image_bytes)
|
|
|
|
|
|
def rembg_available() -> bool:
|
|
try:
|
|
import rembg # noqa: F401
|
|
return True
|
|
except ImportError:
|
|
return False
|
|
|
|
|
|
def gpu_available() -> bool:
|
|
try:
|
|
import torch
|
|
return torch.cuda.is_available()
|
|
except ImportError:
|
|
return False
|