Add LaMa magic eraser, remote provider abstraction, and AI tool infrastructure
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
This commit is contained in:
Claude
+12
-1
@@ -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
|
||||
|
||||
+2
-1
@@ -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")
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user