From 27261c4ef47d2839c942d450d3c062100fb6455d Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 9 Jun 2026 17:42:48 +0000 Subject: [PATCH] Add LaMa magic eraser, remote provider abstraction, and AI tool infrastructure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .env.example | 19 +- backend/app/config.py | 13 +- backend/app/main.py | 3 +- backend/app/routers/ai_tools.py | 281 ++++++++++++ backend/app/services/local_inpaint.py | 82 ++++ backend/app/services/remote_provider.py | 431 ++++++++++++++++++ backend/requirements.txt | 15 +- frontend/src/js/api/capabilities.js | 56 +++ frontend/src/js/config.js | 23 + .../src/js/core/components/provider-badge.js | 61 +++ frontend/src/js/main.js | 4 + frontend/src/js/services/api.js | 118 +++++ frontend/src/js/tools/ai_lama_erase.js | 199 ++++++++ frontend/src/js/tools/ai_smart_inpaint.js | 201 ++++++++ 14 files changed, 1493 insertions(+), 13 deletions(-) create mode 100644 backend/app/routers/ai_tools.py create mode 100644 backend/app/services/local_inpaint.py create mode 100644 backend/app/services/remote_provider.py create mode 100644 frontend/src/js/api/capabilities.js create mode 100644 frontend/src/js/core/components/provider-badge.js create mode 100644 frontend/src/js/tools/ai_lama_erase.js create mode 100644 frontend/src/js/tools/ai_smart_inpaint.js diff --git a/.env.example b/.env.example index eb1b3ab..38d7804 100644 --- a/.env.example +++ b/.env.example @@ -50,10 +50,27 @@ AI_PROVIDER=replicate REPLICATE_API_KEY=r8_PASTE_YOUR_KEY_HERE # ─────────────────────────────────────────────────────────────────────────── -# OPENAI (Alternative - not recommended, lower quality) +# OPENAI (cloud, dall-e-3 / gpt-image-1) # Get key at: https://platform.openai.com/api-keys +# AI_PROVIDER=openai # ─────────────────────────────────────────────────────────────────────────── #OPENAI_API_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx +#OPENAI_MODEL=dall-e-3 + +# ─────────────────────────────────────────────────────────────────────────── +# INVOKEAI (self-hosted, best for Flux/SDXL) +# Run InvokeAI on your local machine or NAS, point URL here. +# AI_PROVIDER=invokeai +# ─────────────────────────────────────────────────────────────────────────── +#INVOKEAI_URL=http://192.168.1.x:9090 +#INVOKEAI_DEFAULT_MODEL=flux-dev + +# ─────────────────────────────────────────────────────────────────────────── +# COMFYUI (self-hosted, workflow JSON API) +# AI_PROVIDER=comfyui +# ─────────────────────────────────────────────────────────────────────────── +#COMFYUI_URL=http://192.168.1.x:8188 +#COMFYUI_DEFAULT_MODEL=v1-5-pruned-emaonly.ckpt # ─────────────────────────────────────────────────────────────────────────── # STABILITY AI (Alternative) diff --git a/backend/app/config.py b/backend/app/config.py index c958edc..1ac4fb4 100644 --- a/backend/app/config.py +++ b/backend/app/config.py @@ -12,13 +12,24 @@ class Settings(BaseSettings): access_token_expire_minutes: int = 30 # AI Provider - ai_provider: str = "mock" # Options: openai, stability, replicate, mock + # Local: blank or "mock" — always available, no config needed + # Remote (set ONE): openai | invokeai | comfyui | replicate | stability + ai_provider: str = "mock" # Provider API Keys openai_api_key: str = "" + openai_model: str = "dall-e-3" stability_api_key: str = "" replicate_api_key: str = "" + # InvokeAI (self-hosted) + invokeai_url: str = "" + invokeai_default_model: str = "flux-dev" + + # ComfyUI (self-hosted) + comfyui_url: str = "" + comfyui_default_model: str = "v1-5-pruned-emaonly.ckpt" + # Model Selection (optional, provider-specific) stability_model: str = "sdxl" # Options: sdxl, sd15, sd21 replicate_model: str = "sdxl-inpaint" # Options: sdxl-inpaint, lama, realistic-vision diff --git a/backend/app/main.py b/backend/app/main.py index bfbb7f2..bf07fe1 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -8,7 +8,7 @@ import os from app.config import settings from app.database import init_db -from app.routers import projects, edits, images, patches, generate, tools +from app.routers import projects, edits, images, patches, generate, tools, ai_tools @asynccontextmanager @@ -41,6 +41,7 @@ app.include_router(images.router) app.include_router(patches.router) app.include_router(generate.router) app.include_router(tools.router) +app.include_router(ai_tools.router) @app.get("/api") diff --git a/backend/app/routers/ai_tools.py b/backend/app/routers/ai_tools.py new file mode 100644 index 0000000..edaa896 --- /dev/null +++ b/backend/app/routers/ai_tools.py @@ -0,0 +1,281 @@ +""" +AI tools router — LaMa inpaint, background removal, remote generation, config. +All endpoints are under /api prefix. +""" + +from fastapi import APIRouter, HTTPException +from pydantic import BaseModel +from typing import Optional +import base64 +import asyncio + +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(): + from app.services.remote_provider import get_remote_provider + provider = get_remote_provider() + if provider is None: + raise HTTPException( + status_code=503, + detail="No remote AI provider configured. Set 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() + 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.post("/generate/txt2img") +async def txt2img(req: Txt2ImgRequest): + """Text-to-image via configured remote provider.""" + provider = _require_remote() + 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() + 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() + 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 ──────────────────────────────────────────────────── + +@router.get("/config") +async def get_config(): + """ + Return capability flags so the frontend can show/hide tools. + Frontend reads this on load. + """ + from app.services.remote_provider import get_remote_provider + from app.config import settings + + remote_caps: list[str] = [] + remote_healthy = False + provider_name = (settings.ai_provider or "").lower() + + if provider_name in ("openai", "invokeai", "comfyui"): + try: + provider = get_remote_provider() + if provider: + remote_caps = provider.capabilities() + remote_healthy = await asyncio.wait_for(provider.health(), timeout=5.0) + except Exception: + remote_healthy = False + + return { + "local": { + "lama": lama_available(), + "rembg": rembg_available(), + "opencv": True, + "gpu_detected": gpu_available(), + }, + "remote": { + "provider": provider_name or None, + "capabilities": remote_caps, + "healthy": remote_healthy, + } + } diff --git a/backend/app/services/local_inpaint.py b/backend/app/services/local_inpaint.py new file mode 100644 index 0000000..cd82997 --- /dev/null +++ b/backend/app/services/local_inpaint.py @@ -0,0 +1,82 @@ +""" +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 diff --git a/backend/app/services/remote_provider.py b/backend/app/services/remote_provider.py new file mode 100644 index 0000000..6e899f7 --- /dev/null +++ b/backend/app/services/remote_provider.py @@ -0,0 +1,431 @@ +""" +Remote AI provider abstraction. +One interface, three drivers: OpenAI, InvokeAI, ComfyUI. +Configure one provider via AI_PROVIDER in .env. +""" + +from abc import ABC, abstractmethod +from typing import Optional +import httpx +import base64 +import asyncio +from io import BytesIO + + +class RemoteAIProvider(ABC): + @abstractmethod + async def inpaint(self, image_bytes: bytes, mask_bytes: bytes, prompt: str, params: dict) -> bytes: ... + @abstractmethod + async def txt2img(self, prompt: str, width: int, height: int, params: dict) -> bytes: ... + @abstractmethod + async def img2img(self, image_bytes: bytes, prompt: str, strength: float, params: dict) -> bytes: ... + @abstractmethod + async def outpaint(self, image_bytes: bytes, direction: str, size: int, prompt: str) -> bytes: ... + @abstractmethod + async def health(self) -> bool: ... + @abstractmethod + def capabilities(self) -> list[str]: ... + + +class OpenAIRemoteProvider(RemoteAIProvider): + """OpenAI image API — gpt-image-1 / dall-e-3.""" + + def __init__(self, api_key: str, model: str = "dall-e-3"): + self.api_key = api_key + self.model = model + self.base_url = "https://api.openai.com/v1" + + def _headers(self): + return {"Authorization": f"Bearer {self.api_key}"} + + async def inpaint(self, image_bytes: bytes, mask_bytes: bytes, prompt: str, params: dict) -> bytes: + async with httpx.AsyncClient(timeout=120.0) as client: + files = { + "image": ("image.png", image_bytes, "image/png"), + "mask": ("mask.png", mask_bytes, "image/png"), + } + data = {"prompt": prompt, "n": "1", "size": "1024x1024"} + r = await client.post(f"{self.base_url}/images/edits", files=files, data=data, headers=self._headers()) + r.raise_for_status() + url = r.json()["data"][0]["url"] + img_r = await client.get(url) + img_r.raise_for_status() + return img_r.content + + async def txt2img(self, prompt: str, width: int, height: int, params: dict) -> bytes: + size = f"{width}x{height}" if f"{width}x{height}" in {"256x256", "512x512", "1024x1024"} else "1024x1024" + async with httpx.AsyncClient(timeout=120.0) as client: + data = {"model": self.model, "prompt": prompt, "n": 1, "size": size} + r = await client.post(f"{self.base_url}/images/generations", json=data, headers=self._headers()) + r.raise_for_status() + url = r.json()["data"][0]["url"] + img_r = await client.get(url) + img_r.raise_for_status() + return img_r.content + + async def img2img(self, image_bytes: bytes, prompt: str, strength: float, params: dict) -> bytes: + # OpenAI doesn't have img2img natively — use edits with blank mask + from PIL import Image + import numpy as np + img = Image.open(BytesIO(image_bytes)).convert("RGBA") + mask = Image.new("RGBA", img.size, (0, 0, 0, 0)) + mask_buf = BytesIO() + mask.save(mask_buf, format="PNG") + return await self.inpaint(image_bytes, mask_buf.getvalue(), prompt, params) + + async def outpaint(self, image_bytes: bytes, direction: str, size: int, prompt: str) -> bytes: + from PIL import Image + img = Image.open(BytesIO(image_bytes)).convert("RGBA") + w, h = img.size + directions = {"left": (size, 0), "right": (size, 0), "top": (0, size), "bottom": (0, size)} + dw, dh = directions.get(direction, (size, 0)) + new_w, new_h = w + dw, h + dh + canvas = Image.new("RGBA", (new_w, new_h), (0, 0, 0, 0)) + offsets = { + "left": (size, 0), "right": (0, 0), "top": (0, size), "bottom": (0, 0) + } + ox, oy = offsets.get(direction, (0, 0)) + canvas.paste(img, (ox, oy)) + # mask: transparent = inpaint + mask = Image.new("L", (new_w, new_h), 0) + # fill the expanded region with white in mask + import numpy as np + mask_arr = np.zeros((new_h, new_w), dtype=np.uint8) + if direction == "left": + mask_arr[:, :size] = 255 + elif direction == "right": + mask_arr[:, w:] = 255 + elif direction == "top": + mask_arr[:size, :] = 255 + else: + mask_arr[h:, :] = 255 + mask = Image.fromarray(mask_arr, "L") + + canvas_rgb = canvas.convert("RGB") + img_buf = BytesIO() + canvas_rgb.save(img_buf, format="PNG") + mask_buf = BytesIO() + mask.save(mask_buf, format="PNG") + return await self.inpaint(img_buf.getvalue(), mask_buf.getvalue(), prompt, {}) + + async def health(self) -> bool: + try: + async with httpx.AsyncClient(timeout=10.0) as client: + r = await client.get(f"{self.base_url}/models", headers=self._headers()) + return r.status_code == 200 + except Exception: + return False + + def capabilities(self) -> list[str]: + return ["inpaint", "txt2img", "img2img", "outpaint"] + + +class InvokeAIProvider(RemoteAIProvider): + """InvokeAI REST API driver — supports Flux, SDXL, SD1.5 and more.""" + + def __init__(self, base_url: str, default_model: str = "flux-dev"): + self.base_url = base_url.rstrip("/") + self.default_model = default_model + + async def _b64(self, data: bytes) -> str: + return base64.b64encode(data).decode() + + async def _upload_image(self, client: httpx.AsyncClient, image_bytes: bytes, category: str = "general") -> str: + """Upload image to InvokeAI and return image_name.""" + files = {"file": ("image.png", image_bytes, "image/png")} + data = {"image_category": category, "is_intermediate": "false"} + r = await client.post(f"{self.base_url}/api/v1/images/upload", files=files, data=data) + r.raise_for_status() + return r.json()["image_name"] + + async def _run_graph(self, client: httpx.AsyncClient, graph: dict) -> bytes: + """Post a graph, poll for completion, return result image bytes.""" + r = await client.post(f"{self.base_url}/api/v1/queue/default/enqueue_batch", + json={"prepend": False, "batch": {"graph": graph, "runs": 1}}) + r.raise_for_status() + batch_id = r.json()["batch"]["batch_id"] + + # Poll queue status + for _ in range(180): + await asyncio.sleep(2) + sr = await client.get(f"{self.base_url}/api/v1/queue/default/status") + sr.raise_for_status() + status = sr.json() + if status.get("queue", {}).get("completed", 0) > 0: + break + if status.get("queue", {}).get("failed", 0) > 0: + raise RuntimeError("InvokeAI graph failed") + + # Fetch latest result image + lr = await client.get(f"{self.base_url}/api/v1/images/?categories=general&limit=1&is_intermediate=false") + lr.raise_for_status() + items = lr.json().get("items", []) + if not items: + raise RuntimeError("No output image from InvokeAI") + + img_name = items[0]["image_name"] + img_r = await client.get(f"{self.base_url}/api/v1/images/i/{img_name}/full") + img_r.raise_for_status() + return img_r.content + + async def inpaint(self, image_bytes: bytes, mask_bytes: bytes, prompt: str, params: dict) -> bytes: + async with httpx.AsyncClient(timeout=300.0) as client: + img_name = await self._upload_image(client, image_bytes) + mask_name = await self._upload_image(client, mask_bytes, "mask") + model = params.get("model", self.default_model) + graph = { + "id": "inpaint_graph", + "nodes": { + "img_node": {"id": "img_node", "type": "image", "image": {"image_name": img_name}}, + "mask_node": {"id": "mask_node", "type": "image", "image": {"image_name": mask_name}}, + "model_node": {"id": "model_node", "type": "main_model_loader", "model": {"model_name": model, "base": "any"}}, + "clip_skip": {"id": "clip_skip", "type": "clip_skip", "skipped_layers": 0}, + "positive": {"id": "positive", "type": "compel", "prompt": prompt}, + "negative": {"id": "negative", "type": "compel", "prompt": params.get("negative_prompt", "")}, + "denoise": { + "id": "denoise", "type": "denoise_latents", + "steps": params.get("steps", 30), + "cfg_scale": params.get("cfg_scale", 7.5), + "denoising_start": 0.0, "denoising_end": 1.0, + "scheduler": "euler", "is_intermediate": False + }, + "vae_loader": {"id": "vae_loader", "type": "vae_loader", "vae_model": {"model_name": model, "base": "any"}}, + "img_to_latents": {"id": "img_to_latents", "type": "i2l"}, + "latents_to_img": {"id": "latents_to_img", "type": "l2i"}, + }, + "edges": [ + {"source": {"node_id": "model_node", "field": "unet"}, "destination": {"node_id": "denoise", "field": "unet"}}, + {"source": {"node_id": "model_node", "field": "clip"}, "destination": {"node_id": "clip_skip", "field": "clip"}}, + {"source": {"node_id": "clip_skip", "field": "clip"}, "destination": {"node_id": "positive", "field": "clip"}}, + {"source": {"node_id": "clip_skip", "field": "clip"}, "destination": {"node_id": "negative", "field": "clip"}}, + {"source": {"node_id": "positive", "field": "conditioning"}, "destination": {"node_id": "denoise", "field": "positive_conditioning"}}, + {"source": {"node_id": "negative", "field": "conditioning"}, "destination": {"node_id": "denoise", "field": "negative_conditioning"}}, + {"source": {"node_id": "img_node", "field": "image"}, "destination": {"node_id": "img_to_latents", "field": "image"}}, + {"source": {"node_id": "vae_loader", "field": "vae"}, "destination": {"node_id": "img_to_latents", "field": "vae"}}, + {"source": {"node_id": "img_to_latents", "field": "latents"}, "destination": {"node_id": "denoise", "field": "latents"}}, + {"source": {"node_id": "mask_node", "field": "image"}, "destination": {"node_id": "denoise", "field": "mask"}}, + {"source": {"node_id": "denoise", "field": "latents"}, "destination": {"node_id": "latents_to_img", "field": "latents"}}, + {"source": {"node_id": "vae_loader", "field": "vae"}, "destination": {"node_id": "latents_to_img", "field": "vae"}}, + ] + } + return await self._run_graph(client, graph) + + async def txt2img(self, prompt: str, width: int, height: int, params: dict) -> bytes: + async with httpx.AsyncClient(timeout=300.0) as client: + model = params.get("model", self.default_model) + graph = { + "id": "txt2img_graph", + "nodes": { + "model_node": {"id": "model_node", "type": "main_model_loader", "model": {"model_name": model, "base": "any"}}, + "clip_skip": {"id": "clip_skip", "type": "clip_skip", "skipped_layers": 0}, + "positive": {"id": "positive", "type": "compel", "prompt": prompt}, + "negative": {"id": "negative", "type": "compel", "prompt": params.get("negative_prompt", "")}, + "noise": {"id": "noise", "type": "noise", "width": width, "height": height, "seed": params.get("seed", 0)}, + "denoise": { + "id": "denoise", "type": "denoise_latents", + "steps": params.get("steps", 30), + "cfg_scale": params.get("cfg_scale", 7.5), + "denoising_start": 0.0, "denoising_end": 1.0, + "scheduler": "euler", + }, + "vae_loader": {"id": "vae_loader", "type": "vae_loader", "vae_model": {"model_name": model, "base": "any"}}, + "latents_to_img": {"id": "latents_to_img", "type": "l2i"}, + }, + "edges": [ + {"source": {"node_id": "model_node", "field": "unet"}, "destination": {"node_id": "denoise", "field": "unet"}}, + {"source": {"node_id": "model_node", "field": "clip"}, "destination": {"node_id": "clip_skip", "field": "clip"}}, + {"source": {"node_id": "clip_skip", "field": "clip"}, "destination": {"node_id": "positive", "field": "clip"}}, + {"source": {"node_id": "clip_skip", "field": "clip"}, "destination": {"node_id": "negative", "field": "clip"}}, + {"source": {"node_id": "positive", "field": "conditioning"}, "destination": {"node_id": "denoise", "field": "positive_conditioning"}}, + {"source": {"node_id": "negative", "field": "conditioning"}, "destination": {"node_id": "denoise", "field": "negative_conditioning"}}, + {"source": {"node_id": "noise", "field": "noise"}, "destination": {"node_id": "denoise", "field": "noise"}}, + {"source": {"node_id": "denoise", "field": "latents"}, "destination": {"node_id": "latents_to_img", "field": "latents"}}, + {"source": {"node_id": "vae_loader", "field": "vae"}, "destination": {"node_id": "latents_to_img", "field": "vae"}}, + ] + } + return await self._run_graph(client, graph) + + async def img2img(self, image_bytes: bytes, prompt: str, strength: float, params: dict) -> bytes: + # Reuse inpaint with a full-white mask at the given strength + from PIL import Image + img = Image.open(BytesIO(image_bytes)) + mask = Image.new("L", img.size, 255) + mask_buf = BytesIO() + mask.save(mask_buf, format="PNG") + p = dict(params) + p.setdefault("denoising_start", 1.0 - strength) + return await self.inpaint(image_bytes, mask_buf.getvalue(), prompt, p) + + async def outpaint(self, image_bytes: bytes, direction: str, size: int, prompt: str) -> bytes: + # Delegate to inpaint with expanded canvas + provider = OpenAIRemoteProvider.__new__(OpenAIRemoteProvider) + return await provider.outpaint(image_bytes, direction, size, prompt) + + async def health(self) -> bool: + try: + async with httpx.AsyncClient(timeout=10.0) as client: + r = await client.get(f"{self.base_url}/api/v1/app/version") + return r.status_code == 200 + except Exception: + return False + + def capabilities(self) -> list[str]: + return ["inpaint", "txt2img", "img2img", "outpaint"] + + +class ComfyUIProvider(RemoteAIProvider): + """ComfyUI workflow JSON API driver.""" + + def __init__(self, base_url: str, default_model: str = "v1-5-pruned-emaonly.ckpt"): + self.base_url = base_url.rstrip("/") + self.default_model = default_model + + async def _upload_image(self, client: httpx.AsyncClient, image_bytes: bytes, name: str = "image.png") -> str: + files = {"image": (name, image_bytes, "image/png")} + data = {"overwrite": "true"} + r = await client.post(f"{self.base_url}/upload/image", files=files, data=data) + r.raise_for_status() + j = r.json() + return j.get("name", name) + + async def _queue_prompt(self, client: httpx.AsyncClient, workflow: dict) -> str: + r = await client.post(f"{self.base_url}/prompt", json={"prompt": workflow}) + r.raise_for_status() + return r.json()["prompt_id"] + + async def _wait_for_result(self, client: httpx.AsyncClient, prompt_id: str) -> bytes: + for _ in range(180): + await asyncio.sleep(2) + r = await client.get(f"{self.base_url}/history/{prompt_id}") + r.raise_for_status() + history = r.json() + if prompt_id in history: + outputs = history[prompt_id].get("outputs", {}) + for node_output in outputs.values(): + for img_info in node_output.get("images", []): + img_r = await client.get( + f"{self.base_url}/view", + params={"filename": img_info["filename"], "subfolder": img_info.get("subfolder", ""), + "type": img_info.get("type", "output")} + ) + img_r.raise_for_status() + return img_r.content + raise RuntimeError("ComfyUI timed out waiting for result") + + async def inpaint(self, image_bytes: bytes, mask_bytes: bytes, prompt: str, params: dict) -> bytes: + model = params.get("model", self.default_model) + async with httpx.AsyncClient(timeout=300.0) as client: + img_name = await self._upload_image(client, image_bytes, "input.png") + mask_name = await self._upload_image(client, mask_bytes, "mask.png") + workflow = { + "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": model}}, + "2": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["1", 1]}}, + "3": {"class_type": "CLIPTextEncode", "inputs": {"text": params.get("negative_prompt", ""), "clip": ["1", 1]}}, + "4": {"class_type": "LoadImage", "inputs": {"image": img_name}}, + "5": {"class_type": "LoadImage", "inputs": {"image": mask_name}}, + "6": {"class_type": "VAEEncode", "inputs": {"pixels": ["4", 0], "vae": ["1", 2]}}, + "7": {"class_type": "KSampler", "inputs": { + "model": ["1", 0], "positive": ["2", 0], "negative": ["3", 0], + "latent_image": ["6", 0], "mask": ["5", 0], + "seed": params.get("seed", 42), "steps": params.get("steps", 20), + "cfg": params.get("cfg_scale", 7.0), "sampler_name": "euler", + "scheduler": "normal", "denoise": params.get("denoise", 1.0) + }}, + "8": {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0], "vae": ["1", 2]}}, + "9": {"class_type": "SaveImage", "inputs": {"images": ["8", 0], "filename_prefix": "api_out"}}, + } + pid = await self._queue_prompt(client, workflow) + return await self._wait_for_result(client, pid) + + async def txt2img(self, prompt: str, width: int, height: int, params: dict) -> bytes: + model = params.get("model", self.default_model) + async with httpx.AsyncClient(timeout=300.0) as client: + workflow = { + "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": model}}, + "2": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["1", 1]}}, + "3": {"class_type": "CLIPTextEncode", "inputs": {"text": params.get("negative_prompt", ""), "clip": ["1", 1]}}, + "4": {"class_type": "EmptyLatentImage", "inputs": {"width": width, "height": height, "batch_size": 1}}, + "5": {"class_type": "KSampler", "inputs": { + "model": ["1", 0], "positive": ["2", 0], "negative": ["3", 0], + "latent_image": ["4", 0], + "seed": params.get("seed", 42), "steps": params.get("steps", 20), + "cfg": params.get("cfg_scale", 7.0), "sampler_name": "euler", + "scheduler": "normal", "denoise": 1.0 + }}, + "6": {"class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["1", 2]}}, + "7": {"class_type": "SaveImage", "inputs": {"images": ["6", 0], "filename_prefix": "api_out"}}, + } + pid = await self._queue_prompt(client, workflow) + return await self._wait_for_result(client, pid) + + async def img2img(self, image_bytes: bytes, prompt: str, strength: float, params: dict) -> bytes: + from PIL import Image + img = Image.open(BytesIO(image_bytes)) + mask = Image.new("L", img.size, 255) + mask_buf = BytesIO() + mask.save(mask_buf, format="PNG") + p = dict(params) + p["denoise"] = strength + return await self.inpaint(image_bytes, mask_buf.getvalue(), prompt, p) + + async def outpaint(self, image_bytes: bytes, direction: str, size: int, prompt: str) -> bytes: + # Build expanded canvas then inpaint with blank mask + from PIL import Image + import numpy as np + img = Image.open(BytesIO(image_bytes)).convert("RGB") + w, h = img.size + dw = size if direction in ("left", "right") else 0 + dh = size if direction in ("top", "bottom") else 0 + canvas = Image.new("RGB", (w + dw, h + dh), (128, 128, 128)) + ox = size if direction == "left" else 0 + oy = size if direction == "top" else 0 + canvas.paste(img, (ox, oy)) + mask_arr = np.zeros((h + dh, w + dw), dtype=np.uint8) + if direction == "left": + mask_arr[:, :size] = 255 + elif direction == "right": + mask_arr[:, w:] = 255 + elif direction == "top": + mask_arr[:size, :] = 255 + else: + mask_arr[h:, :] = 255 + img_buf = BytesIO() + canvas.save(img_buf, format="PNG") + mask_buf = BytesIO() + Image.fromarray(mask_arr, "L").save(mask_buf, format="PNG") + return await self.inpaint(img_buf.getvalue(), mask_buf.getvalue(), prompt, {}) + + async def health(self) -> bool: + try: + async with httpx.AsyncClient(timeout=10.0) as client: + r = await client.get(f"{self.base_url}/system_stats") + return r.status_code == 200 + except Exception: + return False + + def capabilities(self) -> list[str]: + return ["inpaint", "txt2img", "img2img", "outpaint"] + + +def get_remote_provider() -> Optional[RemoteAIProvider]: + """Return the configured remote provider, or None if not configured.""" + from app.config import settings + + provider = (settings.ai_provider or "").lower() + + if provider == "openai": + if not settings.openai_api_key: + return None + return OpenAIRemoteProvider(settings.openai_api_key, settings.openai_model) + + if provider == "invokeai": + if not settings.invokeai_url: + return None + return InvokeAIProvider(settings.invokeai_url, settings.invokeai_default_model) + + if provider == "comfyui": + if not settings.comfyui_url: + return None + return ComfyUIProvider(settings.comfyui_url, settings.comfyui_default_model) + + return None diff --git a/backend/requirements.txt b/backend/requirements.txt index 9649ea5..28e4923 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -12,19 +12,14 @@ httpx==0.26.0 pydantic==2.5.3 pydantic-settings==2.1.0 email-validator==2.1.0 -opencv-python-headless==4.9.0.80 +opencv-python-headless>=4.10.0 # SAM (Segment Anything) for smart object selection - runs locally, no API needed torch==2.1.2 torchvision==0.16.2 segment-anything @ git+https://github.com/facebookresearch/segment-anything.git -# Note: Using OpenCV DNN instead of onnxruntime for U2Net -# (onnxruntime has executable stack issues in some Docker environments) +# Local AI inpainting — LaMa model (auto GPU/CPU, no API key needed) +simple-lama-inpainting -# NOTE: rembg (background removal) disabled due to dependency conflicts -# rembg>=2.0.70 requires: -# - scikit-image>=0.26.0 which requires numpy>=2.0 -# - Pillow>=12.1.0 -# But opencv-python-headless 4.9.0.80 requires numpy<2.0 -# To enable rembg, need to update opencv-python-headless to 4.10+ (numpy 2.x compatible) -# and update all dependent packages accordingly +# Background removal — rembg enabled now that opencv 4.10+ supports numpy 2.x +rembg[gpu] diff --git a/frontend/src/js/api/capabilities.js b/frontend/src/js/api/capabilities.js new file mode 100644 index 0000000..feb6307 --- /dev/null +++ b/frontend/src/js/api/capabilities.js @@ -0,0 +1,56 @@ +/** + * Backend capabilities singleton. + * Fetched once on load from GET /api/config. + * Tools use this to decide whether to show, grey out, or show tooltips. + * + * Shape: + * { + * local: { lama, rembg, opencv, gpu_detected }, + * remote: { provider, capabilities: string[], healthy } + * } + */ + +import apiService from '../services/api.js'; + +const DEFAULT_CAPS = { + local: { lama: false, rembg: false, opencv: true, gpu_detected: false }, + remote: { provider: null, capabilities: [], healthy: false }, +}; + +let _caps = null; +let _fetchPromise = null; + +/** + * Return capabilities (fetched lazily, cached thereafter). + * Always resolves — falls back to DEFAULT_CAPS on network error. + */ +export async function getCapabilities() { + if (_caps) return _caps; + if (!_fetchPromise) { + _fetchPromise = apiService.getConfig() + .then(data => { _caps = data || DEFAULT_CAPS; return _caps; }) + .catch(() => { _caps = DEFAULT_CAPS; return _caps; }); + } + return _fetchPromise; +} + +/** + * Synchronous check — returns cached value or DEFAULT_CAPS if not yet loaded. + */ +export function getCachedCapabilities() { + return _caps || DEFAULT_CAPS; +} + +/** + * True if the remote provider is configured and healthy. + */ +export function hasRemote() { + return !!(_caps?.remote?.healthy); +} + +/** + * Kick off the fetch immediately at module load time so it's ready when tools need it. + */ +getCapabilities(); + +export default { getCapabilities, getCachedCapabilities, hasRemote }; diff --git a/frontend/src/js/config.js b/frontend/src/js/config.js index 077c1f7..a2756fc 100644 --- a/frontend/src/js/config.js +++ b/frontend/src/js/config.js @@ -110,6 +110,29 @@ config.TOOLS = [ on_activate: 'on_activate', attributes: {}, }, + { + name: 'ai_lama_erase', + title: 'AI Magic Erase (LaMa) - Paint over to erase', + attributes: { + size: { + value: 30, + min: 5, + max: 200, + }, + }, + }, + { + name: 'ai_smart_inpaint', + title: 'AI Smart Inpaint - Paint + describe replacement', + on_activate: 'on_activate', + attributes: { + size: { + value: 30, + min: 5, + max: 200, + }, + }, + }, { name: 'magic_wand', title: 'Magic Wand (Color Select)', diff --git a/frontend/src/js/core/components/provider-badge.js b/frontend/src/js/core/components/provider-badge.js new file mode 100644 index 0000000..0af35f2 --- /dev/null +++ b/frontend/src/js/core/components/provider-badge.js @@ -0,0 +1,61 @@ +/** + * ProviderBadge — small DOM element showing the active AI provider. + * Inserted into the toolbar footer on app load. + * + * Green = remote provider healthy + * Yellow = provider configured but unhealthy/unreachable + * Grey = local only (LaMa + OpenCV) + */ + +import { getCapabilities } from '../../api/capabilities.js'; + +export async function mountProviderBadge(container) { + var caps = await getCapabilities(); + + var badge = document.createElement('div'); + badge.id = 'provider-badge'; + badge.style.cssText = [ + 'display:inline-flex', 'align-items:center', 'gap:5px', + 'padding:3px 8px', 'border-radius:10px', + 'font-size:11px', 'font-family:sans-serif', + 'cursor:default', 'user-select:none', + 'margin:4px', 'opacity:0.85', + ].join(';'); + + var dot = document.createElement('span'); + dot.style.cssText = 'width:7px;height:7px;border-radius:50%;display:inline-block;'; + + var label = document.createElement('span'); + + var remote = caps.remote || {}; + var local = caps.local || {}; + + if (remote.provider && remote.healthy) { + dot.style.background = '#44cc44'; + badge.style.background = '#1a2a1a'; + badge.style.color = '#aaffaa'; + label.textContent = remote.provider + (local.gpu_detected ? ' · GPU' : ' · CPU'); + badge.title = 'Remote provider: ' + remote.provider + '\nCapabilities: ' + (remote.capabilities || []).join(', '); + } else if (remote.provider && !remote.healthy) { + dot.style.background = '#ffaa00'; + badge.style.background = '#2a2000'; + badge.style.color = '#ffdd88'; + label.textContent = remote.provider + ' (offline)'; + badge.title = remote.provider + ' is configured but not reachable. Check your .env URL.'; + } else { + dot.style.background = '#888888'; + badge.style.background = '#1a1a1a'; + badge.style.color = '#aaaaaa'; + label.textContent = 'Local' + (local.lama ? ' · LaMa' : '') + (local.gpu_detected ? ' · GPU' : ''); + badge.title = 'Local only. Set AI_PROVIDER in .env to enable generative tools.'; + } + + badge.appendChild(dot); + badge.appendChild(label); + + if (container) { + container.appendChild(badge); + } + + return badge; +} diff --git a/frontend/src/js/main.js b/frontend/src/js/main.js index 9d3eac0..ded0000 100644 --- a/frontend/src/js/main.js +++ b/frontend/src/js/main.js @@ -23,6 +23,7 @@ import Base_search_class from './core/base-search.js'; import File_open_class from './modules/file/open.js'; import File_save_class from './modules/file/save.js'; import * as Actions from './actions/index.js'; +import { mountProviderBadge } from './core/components/provider-badge.js'; window.addEventListener('load', function (e) { // Initiate app @@ -54,4 +55,7 @@ window.addEventListener('load', function (e) { // Render all GUI.init(); Layers.init(); + + // Mount provider badge in the tools panel footer + mountProviderBadge(document.getElementById('tools_container') || document.body); }, false); diff --git a/frontend/src/js/services/api.js b/frontend/src/js/services/api.js index ee4abb3..f853673 100644 --- a/frontend/src/js/services/api.js +++ b/frontend/src/js/services/api.js @@ -95,6 +95,124 @@ class ApiService { return response.json(); } + /** + * AI erase using LaMa (local, no API key needed) + * @param {string} imageData - Base64 encoded image + * @param {string} maskData - Base64 encoded mask (white = erase) + * @returns {Promise<{result: string, method: string}>} + */ + async erase(imageData, maskData) { + const response = await fetch(`${this.baseUrl}/api/erase`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ image: imageData, mask: maskData }), + }); + if (!response.ok) { + const error = await response.json().catch(() => ({ detail: 'Unknown error' })); + throw new Error(error.detail || `Erase request failed: ${response.status}`); + } + return response.json(); + } + + /** + * Text-to-image via remote provider + * @param {string} prompt + * @param {Object} options - width, height, negativePrompt, steps, cfgScale, model + * @returns {Promise<{result: string}>} + */ + async textToImage(prompt, options = {}) { + const response = await fetch(`${this.baseUrl}/api/generate/txt2img`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + prompt, + width: options.width || 1024, + height: options.height || 1024, + negative_prompt: options.negativePrompt || '', + steps: options.steps || 30, + cfg_scale: options.cfgScale || 7.5, + model: options.model || null, + seed: options.seed || 0, + }), + }); + if (!response.ok) { + const error = await response.json().catch(() => ({ detail: 'Unknown error' })); + throw new Error(error.detail || `Text-to-image failed: ${response.status}`); + } + return response.json(); + } + + /** + * Image-to-image via remote provider + * @param {string} imageData - Base64 encoded image + * @param {string} prompt + * @param {Object} options + * @returns {Promise<{result: string}>} + */ + async imageToImage(imageData, prompt, options = {}) { + const response = await fetch(`${this.baseUrl}/api/generate/img2img`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + image: imageData, + prompt, + strength: options.strength || 0.75, + negative_prompt: options.negativePrompt || '', + steps: options.steps || 30, + cfg_scale: options.cfgScale || 7.5, + model: options.model || null, + }), + }); + if (!response.ok) { + const error = await response.json().catch(() => ({ detail: 'Unknown error' })); + throw new Error(error.detail || `Image-to-image failed: ${response.status}`); + } + return response.json(); + } + + /** + * Inpaint with prompt via remote provider + * @param {string} imageData - Base64 + * @param {string} maskData - Base64 + * @param {string} prompt + * @param {Object} options + * @returns {Promise<{result: string}>} + */ + async remoteInpaint(imageData, maskData, prompt, options = {}) { + const response = await fetch(`${this.baseUrl}/api/inpaint/remote`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + image: imageData, + mask: maskData, + prompt, + negative_prompt: options.negativePrompt || '', + steps: options.steps || 30, + cfg_scale: options.cfgScale || 7.5, + model: options.model || null, + }), + }); + if (!response.ok) { + const error = await response.json().catch(() => ({ detail: 'Unknown error' })); + throw new Error(error.detail || `Remote inpaint failed: ${response.status}`); + } + return response.json(); + } + + /** + * Fetch backend capabilities (local tools available, remote provider status). + * @returns {Promise} + */ + async getConfig() { + try { + const response = await fetch(`${this.baseUrl}/api/config`); + if (!response.ok) return null; + return response.json(); + } catch { + return null; + } + } + /** * Health check for the backend * @returns {Promise} diff --git a/frontend/src/js/tools/ai_lama_erase.js b/frontend/src/js/tools/ai_lama_erase.js new file mode 100644 index 0000000..e4616b6 --- /dev/null +++ b/frontend/src/js/tools/ai_lama_erase.js @@ -0,0 +1,199 @@ +/** + * AI Magic Eraser — paint a mask with a brush, send to LaMa backend, apply result. + * Works locally (no API key). GPU auto-detected; CPU fallback always available. + * + * Workflow: + * 1. User paints over the object to erase (red overlay shows the mask) + * 2. On mouseup, POST image + mask to /api/erase + * 3. Result replaces the current layer canvas + * + * Registered as tool name: "ai_lama_erase" + */ + +import app from './../app.js'; +import config from './../config.js'; +import Base_tools_class from './../core/base-tools.js'; +import Base_layers_class from './../core/base-layers.js'; +import Helper_class from './../libs/helpers.js'; +import alertify from './../../../node_modules/alertifyjs/build/alertify.min.js'; +import apiService from './../services/api.js'; + +class Ai_lama_erase_class extends Base_tools_class { + + constructor(ctx) { + super(); + this.Base_layers = new Base_layers_class(); + this.Helper = new Helper_class(); + this.ctx = ctx; + this.name = 'ai_lama_erase'; + + this.isDrawing = false; + this.isProcessing = false; + + // Off-screen canvas used to accumulate the painted mask + this.maskCanvas = null; + this.maskCtx = null; + } + + load() { + var _this = this; + document.addEventListener('mousedown', function (e) { _this.mousedown(e); }); + document.addEventListener('mousemove', function (e) { _this.mousemove(e); }); + document.addEventListener('mouseup', function (e) { _this.mouseup(e); }); + document.addEventListener('touchstart', function (e) { _this.mousedown(e); }, { passive: false }); + document.addEventListener('touchmove', function (e) { _this.mousemove(e); }, { passive: false }); + document.addEventListener('touchend', function (e) { _this.mouseup(e); }); + } + + mousedown(e) { + var mouse = this.get_mouse_info(e); + if (!mouse.click_valid) return; + if (config.TOOL.name !== this.name) return; + if (this.isProcessing) return; + + if (config.layer.type !== 'image') { + alertify.error('This layer must contain an image.'); + return; + } + + this._initMask(); + this.isDrawing = true; + this._paint(mouse); + } + + mousemove(e) { + if (!this.isDrawing) return; + if (config.TOOL.name !== this.name) return; + var mouse = this.get_mouse_info(e); + this._paint(mouse); + } + + mouseup(e) { + if (!this.isDrawing) return; + this.isDrawing = false; + if (config.TOOL.name !== this.name) return; + this._applyErase(); + } + + // ── Private ────────────────────────────────────────────────────────────── + + _initMask() { + var w = config.layer.width_original; + var h = config.layer.height_original; + + if (!this.maskCanvas || this.maskCanvas.width !== w || this.maskCanvas.height !== h) { + this.maskCanvas = document.createElement('canvas'); + this.maskCanvas.width = w; + this.maskCanvas.height = h; + this.maskCtx = this.maskCanvas.getContext('2d'); + } + this.maskCtx.clearRect(0, 0, w, h); + } + + _paint(mouse) { + var params = this.getParams(); + var size = params.size || 30; + + // Map screen coords → layer-original coords + var lx = Math.round(this.adaptSize(Math.round(mouse.x) - config.layer.x, 'width')); + var ly = Math.round(this.adaptSize(Math.round(mouse.y) - config.layer.y, 'height')); + + this.maskCtx.beginPath(); + this.maskCtx.arc(lx, ly, size / 2, 0, Math.PI * 2); + this.maskCtx.fillStyle = '#ffffff'; + this.maskCtx.fill(); + + // Show red overlay on screen so user can see the painted area + this._renderOverlay(lx, ly, size); + } + + _renderOverlay(lx, ly, size) { + // Draw a translucent red circle on the main canvas for visual feedback + var scale = config.ZOOM / 100; + var sx = config.layer.x * scale + lx * scale; + var sy = config.layer.y * scale + ly * scale; + var sRadius = (size / 2) * scale; + + var mainCtx = document.getElementById('canvas_temp') + ? document.getElementById('canvas_temp').getContext('2d') + : null; + if (!mainCtx) return; + + mainCtx.save(); + mainCtx.beginPath(); + mainCtx.arc(sx, sy, sRadius, 0, Math.PI * 2); + mainCtx.fillStyle = 'rgba(255, 60, 60, 0.4)'; + mainCtx.fill(); + mainCtx.restore(); + } + + async _applyErase() { + if (this.isProcessing) return; + + // Check if any mask pixels were painted + var maskData = this.maskCtx.getImageData( + 0, 0, this.maskCanvas.width, this.maskCanvas.height + ); + var hasPixels = maskData.data.some((v, i) => i % 4 === 3 && v > 0); + if (!hasPixels) return; + + this.isProcessing = true; + alertify.message('AI erasing... please wait', 0); + + try { + // Get current layer as PNG base64 + var layerCanvas = document.createElement('canvas'); + layerCanvas.width = config.layer.width_original; + layerCanvas.height = config.layer.height_original; + var lctx = layerCanvas.getContext('2d'); + lctx.drawImage(config.layer.link, 0, 0); + var imageB64 = layerCanvas.toDataURL('image/png').split(',')[1]; + + // Get mask as PNG base64 + var maskB64 = this.maskCanvas.toDataURL('image/png').split(',')[1]; + + // Call backend + var result = await apiService.erase(imageB64, maskB64); + + // Apply result back to layer + var img = new Image(); + img.onload = () => { + var resultCanvas = document.createElement('canvas'); + resultCanvas.width = config.layer.width_original; + resultCanvas.height = config.layer.height_original; + resultCanvas.getContext('2d').drawImage(img, 0, 0); + + app.State.do_action( + new app.Actions.Bundle_action('ai_lama_erase', 'AI Erase', [ + new app.Actions.Update_layer_image_action(resultCanvas) + ]) + ); + + alertify.dismissAll(); + alertify.success('Erased! (' + result.method + ')'); + this.isProcessing = false; + this._clearOverlay(); + }; + img.onerror = () => { + alertify.dismissAll(); + alertify.error('Failed to load result image.'); + this.isProcessing = false; + }; + img.src = 'data:image/png;base64,' + result.result; + + } catch (err) { + alertify.dismissAll(); + alertify.error('AI erase failed: ' + (err.message || err)); + this.isProcessing = false; + } + } + + _clearOverlay() { + var canvas = document.getElementById('canvas_temp'); + if (canvas) { + canvas.getContext('2d').clearRect(0, 0, canvas.width, canvas.height); + } + } +} + +export default Ai_lama_erase_class; diff --git a/frontend/src/js/tools/ai_smart_inpaint.js b/frontend/src/js/tools/ai_smart_inpaint.js new file mode 100644 index 0000000..61f6322 --- /dev/null +++ b/frontend/src/js/tools/ai_smart_inpaint.js @@ -0,0 +1,201 @@ +/** + * AI Smart Inpaint — paint a mask, enter a prompt, choose Fast (LaMa) or Quality (remote). + * + * Fast mode: /api/erase — LaMa local, no API key, seconds + * Quality mode: /api/inpaint/remote — InvokeAI / ComfyUI / OpenAI, requires configured provider + * + * Registered as tool name: "ai_smart_inpaint" + */ + +import app from './../app.js'; +import config from './../config.js'; +import Base_tools_class from './../core/base-tools.js'; +import Base_layers_class from './../core/base-layers.js'; +import Helper_class from './../libs/helpers.js'; +import Dialog_class from './../libs/popup.js'; +import alertify from './../../../node_modules/alertifyjs/build/alertify.min.js'; +import apiService from './../services/api.js'; +import { getCapabilities } from './../api/capabilities.js'; + +class Ai_smart_inpaint_class extends Base_tools_class { + + constructor(ctx) { + super(); + this.Base_layers = new Base_layers_class(); + this.Helper = new Helper_class(); + this.POP = new Dialog_class(); + this.ctx = ctx; + this.name = 'ai_smart_inpaint'; + + this.isDrawing = false; + this.isProcessing = false; + this.maskCanvas = null; + this.maskCtx = null; + } + + load() { + var _this = this; + document.addEventListener('mousedown', function (e) { _this.mousedown(e); }); + document.addEventListener('mousemove', function (e) { _this.mousemove(e); }); + document.addEventListener('mouseup', function (e) { _this.mouseup(e); }); + document.addEventListener('touchstart', function (e) { _this.mousedown(e); }, { passive: false }); + document.addEventListener('touchmove', function (e) { _this.mousemove(e); }, { passive: false }); + document.addEventListener('touchend', function (e) { _this.mouseup(e); }); + } + + on_activate() { + // Nothing on activate — tool is drag-to-paint, then dialog on mouseup + } + + mousedown(e) { + var mouse = this.get_mouse_info(e); + if (!mouse.click_valid) return; + if (config.TOOL.name !== this.name) return; + if (this.isProcessing) return; + if (config.layer.type !== 'image') { + alertify.error('This layer must contain an image.'); + return; + } + this._initMask(); + this.isDrawing = true; + this._paint(mouse); + } + + mousemove(e) { + if (!this.isDrawing) return; + if (config.TOOL.name !== this.name) return; + this._paint(this.get_mouse_info(e)); + } + + mouseup(e) { + if (!this.isDrawing) return; + this.isDrawing = false; + if (config.TOOL.name !== this.name) return; + + var maskData = this.maskCtx.getImageData( + 0, 0, this.maskCanvas.width, this.maskCanvas.height + ); + if (!maskData.data.some((v, i) => i % 4 === 3 && v > 0)) return; + + this._showDialog(); + } + + // ── Private ────────────────────────────────────────────────────────────── + + _initMask() { + var w = config.layer.width_original; + var h = config.layer.height_original; + if (!this.maskCanvas || this.maskCanvas.width !== w || this.maskCanvas.height !== h) { + this.maskCanvas = document.createElement('canvas'); + this.maskCanvas.width = w; + this.maskCanvas.height = h; + this.maskCtx = this.maskCanvas.getContext('2d'); + } + this.maskCtx.clearRect(0, 0, w, h); + } + + _paint(mouse) { + var params = this.getParams(); + var size = params.size || 30; + var lx = Math.round(this.adaptSize(Math.round(mouse.x) - config.layer.x, 'width')); + var ly = Math.round(this.adaptSize(Math.round(mouse.y) - config.layer.y, 'height')); + this.maskCtx.beginPath(); + this.maskCtx.arc(lx, ly, size / 2, 0, Math.PI * 2); + this.maskCtx.fillStyle = '#ffffff'; + this.maskCtx.fill(); + } + + async _showDialog() { + var caps = await getCapabilities(); + var hasRemote = caps.remote && caps.remote.healthy; + + var _this = this; + + var settings = { + title: 'AI Smart Inpaint', + params: [ + { + name: 'quality', + title: 'Mode:', + value: 'fast', + values: hasRemote ? ['fast', 'quality'] : ['fast'], + note: hasRemote ? 'Fast = LaMa (local). Quality = remote AI + prompt.' : 'Quality mode requires a remote provider (InvokeAI / ComfyUI / OpenAI).', + }, + { + name: 'prompt', + title: 'What to put here (Quality mode only):', + type: 'textarea', + value: '', + placeholder: "e.g. 'lush green grass', 'wooden table surface', 'clear blue sky'", + }, + { + name: 'negative_prompt', + title: 'Avoid (optional):', + value: '', + placeholder: 'blurry, distorted', + }, + ], + on_load: function (params, popup) {}, + on_finish: function (params) { + _this._runInpaint(params.quality, params.prompt, params.negative_prompt); + }, + }; + + this.POP.show(settings); + } + + async _runInpaint(quality, prompt, negativePrompt) { + if (this.isProcessing) return; + this.isProcessing = true; + + var modeLabel = quality === 'quality' ? 'Quality (remote)' : 'Fast (LaMa)'; + alertify.message('Inpainting (' + modeLabel + ')... please wait', 0); + + try { + var layerCanvas = document.createElement('canvas'); + layerCanvas.width = config.layer.width_original; + layerCanvas.height = config.layer.height_original; + layerCanvas.getContext('2d').drawImage(config.layer.link, 0, 0); + var imageB64 = layerCanvas.toDataURL('image/png').split(',')[1]; + var maskB64 = this.maskCanvas.toDataURL('image/png').split(',')[1]; + + var result; + if (quality === 'quality') { + result = await apiService.remoteInpaint(imageB64, maskB64, prompt || 'fill naturally', { negativePrompt }); + } else { + result = await apiService.erase(imageB64, maskB64); + } + + var img = new Image(); + img.onload = () => { + var resultCanvas = document.createElement('canvas'); + resultCanvas.width = config.layer.width_original; + resultCanvas.height = config.layer.height_original; + resultCanvas.getContext('2d').drawImage(img, 0, 0); + + app.State.do_action( + new app.Actions.Bundle_action('ai_smart_inpaint', 'AI Smart Inpaint', [ + new app.Actions.Update_layer_image_action(resultCanvas) + ]) + ); + + alertify.dismissAll(); + alertify.success('Done!'); + this.isProcessing = false; + }; + img.onerror = () => { + alertify.dismissAll(); + alertify.error('Failed to load result.'); + this.isProcessing = false; + }; + img.src = 'data:image/png;base64,' + result.result; + + } catch (err) { + alertify.dismissAll(); + alertify.error('Inpaint failed: ' + (err.message || err)); + this.isProcessing = false; + } + } +} + +export default Ai_smart_inpaint_class;