paintplus: vendor the app source and rename from EditmaskwithAI
Bring the full EditmaskwithAI application into the repo under paintplus/ (429 files) so the service is self-contained — the installer copies the vendored source to ~/docker/paintplus/src instead of cloning at runtime. Rename to PaintPlus (service + branding; app logic untouched): - services/editmaskwithai.sh -> services/paintplus.sh (register_service paintplus, install_paintplus, ~/docker/paintplus, Caddy paintplus:8000, Authelia option preserved) - container names -> paintplus across docker-compose*.yml; dev network -> paintplus-network - browser <title> -> "PaintPlus - AI Image Editor"; README heading -> PaintPlus with upstream provenance note - README utilities table: editmaskwithai -> paintplus Backend/frontend code (help strings referencing the old container name, the ai_photo_edit.db filename) is intentionally left as-is to avoid touching application logic. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Nb2vJ8W7bHKx1JXVvpCraH
This commit is contained in:
Claude
@@ -0,0 +1,79 @@
|
||||
from pydantic_settings import BaseSettings
|
||||
from typing import List
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
# Database
|
||||
database_url: str = "sqlite:///./data/ai_photo_edit.db"
|
||||
|
||||
# Security
|
||||
secret_key: str = "your-secret-key-change-in-production"
|
||||
algorithm: str = "HS256"
|
||||
access_token_expire_minutes: int = 30
|
||||
|
||||
# AI Provider
|
||||
# Local: blank or "mock" — always available, no config needed
|
||||
# Remote default (used for any operation without a specific override):
|
||||
# openai | invokeai | comfyui | replicate | stability
|
||||
ai_provider: str = "mock"
|
||||
|
||||
# Per-operation provider overrides — blank means use ai_provider default.
|
||||
# Operations: inpaint, txt2img, img2img, outpaint
|
||||
# Example: AI_PROVIDER_TXT2IMG=openai (use OpenAI for text-to-image only)
|
||||
ai_provider_inpaint: str = "" # remote inpaint / replace selection
|
||||
ai_provider_txt2img: str = "" # text-to-image
|
||||
ai_provider_img2img: str = "" # image-to-image
|
||||
ai_provider_outpaint: str = "" # expand canvas
|
||||
|
||||
# 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
|
||||
|
||||
# Allow per-edit model override
|
||||
allow_model_override: bool = True
|
||||
|
||||
# Remove Background — preferred local model when request.model="auto"
|
||||
# Options: ben2 (default, best for clean cutouts/hair), birefnet-hr (best
|
||||
# for high-res/print work), u2net (lightweight, smallest download)
|
||||
bg_removal_model: str = "ben2"
|
||||
|
||||
# Local GPU diffusion (AI_PROVIDER=local_gpu)
|
||||
auto_download_models: bool = True # download HF models on first use
|
||||
local_gpu_max_pipelines: int = 2 # max diffusion pipelines kept in GPU memory
|
||||
hf_token: str = "" # HuggingFace token (only needed for gated models)
|
||||
# Override auto-selected models per operation (leave blank = auto-pick by VRAM tier)
|
||||
hf_model_inpaint: str = ""
|
||||
hf_model_txt2img: str = ""
|
||||
hf_model_img2img: str = ""
|
||||
|
||||
# File Storage
|
||||
data_dir: str = "./data"
|
||||
max_upload_size_mb: int = 50
|
||||
|
||||
# CORS
|
||||
cors_origins: str = "http://localhost:3000,http://localhost:5173"
|
||||
|
||||
@property
|
||||
def cors_origins_list(self) -> List[str]:
|
||||
return [origin.strip() for origin in self.cors_origins.split(",")]
|
||||
|
||||
class Config:
|
||||
env_file = ".env"
|
||||
case_sensitive = False
|
||||
|
||||
|
||||
settings = Settings()
|
||||
@@ -0,0 +1,30 @@
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.ext.declarative import declarative_base
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from app.config import settings
|
||||
import os
|
||||
|
||||
# Ensure data directory exists
|
||||
os.makedirs("./data", exist_ok=True)
|
||||
|
||||
engine = create_engine(
|
||||
settings.database_url,
|
||||
connect_args={"check_same_thread": False} # Needed for SQLite
|
||||
)
|
||||
|
||||
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
||||
|
||||
Base = declarative_base()
|
||||
|
||||
|
||||
def get_db():
|
||||
db = SessionLocal()
|
||||
try:
|
||||
yield db
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def init_db():
|
||||
"""Initialize database tables"""
|
||||
Base.metadata.create_all(bind=engine)
|
||||
@@ -0,0 +1,162 @@
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.responses import FileResponse, HTMLResponse
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from app.config import settings
|
||||
from app.database import init_db
|
||||
from app.routers import projects, edits, images, patches, generate, tools, ai_tools, print_tools
|
||||
from app.routers import gpu_status
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
"""Initialize database on startup; auto-install Real-ESRGAN NCNN in background."""
|
||||
init_db()
|
||||
# Kick off NCNN install in background if no AI upscaler detected
|
||||
from app.services.upscale import probe_upscale_capabilities, ensure_ncnn_installed
|
||||
caps = probe_upscale_capabilities()
|
||||
if not caps["realesrgan_pytorch"] and not caps["realesrgan_ncnn"]:
|
||||
asyncio.create_task(ensure_ncnn_installed())
|
||||
# Pre-download SAM model in background so first click is fast
|
||||
from app.services.sam_service import ensure_sam_installed
|
||||
asyncio.create_task(ensure_sam_installed())
|
||||
|
||||
# If local GPU provider is active, log GPU info at startup
|
||||
if settings.ai_provider.lower() == "local_gpu" or any(
|
||||
v.lower() == "local_gpu"
|
||||
for v in [
|
||||
settings.ai_provider_inpaint,
|
||||
settings.ai_provider_txt2img,
|
||||
settings.ai_provider_img2img,
|
||||
settings.ai_provider_outpaint,
|
||||
]
|
||||
if v
|
||||
):
|
||||
from app.services.gpu_detect import get_cached_gpu_info
|
||||
info = get_cached_gpu_info()
|
||||
cc_str = f" | CC={info.compute_capability}" if info.compute_capability else ""
|
||||
print(
|
||||
f"[gpu] {info.device_name} | {info.vram_total_gb:.1f} GB{cc_str} | "
|
||||
f"tier={info.tier} | fp16={info.fp16}"
|
||||
)
|
||||
for w in info.warnings:
|
||||
print(f"[gpu] ⚠ {w}")
|
||||
if settings.auto_download_models:
|
||||
# Download model weight files to disk cache in background so first
|
||||
# user request loads from local disk instead of the internet.
|
||||
from app.services.local_diffusion import prefetch_model_files
|
||||
asyncio.create_task(prefetch_model_files())
|
||||
|
||||
yield
|
||||
|
||||
|
||||
app = FastAPI(
|
||||
title="AI Photo Edit API",
|
||||
description="API for AI-powered photo editing with mask-scoped regeneration",
|
||||
version="1.0.0",
|
||||
lifespan=lifespan
|
||||
)
|
||||
|
||||
# Configure CORS
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=settings.cors_origins_list,
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# Include routers
|
||||
app.include_router(projects.router)
|
||||
app.include_router(edits.router)
|
||||
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.include_router(print_tools.router)
|
||||
app.include_router(gpu_status.router)
|
||||
|
||||
|
||||
@app.get("/api")
|
||||
def api_root():
|
||||
"""API info endpoint"""
|
||||
return {
|
||||
"name": "AI Photo Edit API",
|
||||
"version": "1.0.0",
|
||||
"status": "running"
|
||||
}
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health():
|
||||
"""Health check endpoint"""
|
||||
return {"status": "healthy"}
|
||||
|
||||
|
||||
# Static files directory
|
||||
STATIC_DIR = Path("/app/static")
|
||||
|
||||
|
||||
class NoCacheStaticFiles(StaticFiles):
|
||||
"""webpack outputs a fixed 'bundle.js' filename (no content hash), so
|
||||
browsers can keep serving a stale cached copy after a rebuild unless
|
||||
forced to revalidate on every request."""
|
||||
|
||||
def file_response(self, *args, **kwargs):
|
||||
response = super().file_response(*args, **kwargs)
|
||||
response.headers["Cache-Control"] = "no-cache"
|
||||
return response
|
||||
|
||||
|
||||
# Serve static assets - mount subdirectories if they exist
|
||||
if STATIC_DIR.exists():
|
||||
# React-style assets folder
|
||||
if (STATIC_DIR / "assets").exists():
|
||||
app.mount("/assets", StaticFiles(directory=STATIC_DIR / "assets"), name="assets")
|
||||
# miniPaint dist folder (webpack bundle) - no-cache so code updates are picked up immediately
|
||||
if (STATIC_DIR / "dist").exists():
|
||||
app.mount("/dist", NoCacheStaticFiles(directory=STATIC_DIR / "dist"), name="dist")
|
||||
# miniPaint images folder
|
||||
if (STATIC_DIR / "images").exists():
|
||||
app.mount("/images", StaticFiles(directory=STATIC_DIR / "images"), name="images")
|
||||
# miniPaint CSS folder - no-cache, same reasoning as /dist
|
||||
if (STATIC_DIR / "src").exists():
|
||||
app.mount("/src", NoCacheStaticFiles(directory=STATIC_DIR / "src"), name="src")
|
||||
|
||||
|
||||
@app.get("/", response_class=HTMLResponse)
|
||||
async def serve_spa():
|
||||
"""Serve miniPaint index.html"""
|
||||
index_path = STATIC_DIR / "index.html"
|
||||
if index_path.exists():
|
||||
return FileResponse(index_path, headers={"Cache-Control": "no-cache"})
|
||||
return HTMLResponse("<h1>Frontend not built. Run npm build in frontend/</h1>")
|
||||
|
||||
|
||||
@app.get("/{full_path:path}")
|
||||
async def serve_spa_routes(request: Request, full_path: str):
|
||||
"""
|
||||
Catch-all route for serving static files.
|
||||
Serves static files if they exist, otherwise returns index.html.
|
||||
"""
|
||||
# Don't catch API routes
|
||||
if full_path.startswith(("projects", "edits", "patches", "tools", "generate", "health", "docs", "openapi.json", "api")):
|
||||
return {"detail": "Not Found"}
|
||||
|
||||
# Check if it's a static file
|
||||
static_file = STATIC_DIR / full_path
|
||||
if static_file.exists() and static_file.is_file():
|
||||
return FileResponse(static_file)
|
||||
|
||||
# Otherwise serve index.html
|
||||
index_path = STATIC_DIR / "index.html"
|
||||
if index_path.exists():
|
||||
return FileResponse(index_path, headers={"Cache-Control": "no-cache"})
|
||||
|
||||
return HTMLResponse("<h1>Frontend not built</h1>", status_code=404)
|
||||
@@ -0,0 +1,6 @@
|
||||
from app.models.user import User
|
||||
from app.models.project import Project
|
||||
from app.models.edit import Edit
|
||||
from app.models.patch import Patch
|
||||
|
||||
__all__ = ["User", "Project", "Edit", "Patch"]
|
||||
@@ -0,0 +1,23 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Text
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime
|
||||
from app.database import Base
|
||||
|
||||
|
||||
class Edit(Base):
|
||||
__tablename__ = "edits"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
project_id = Column(Integer, ForeignKey("projects.id"), nullable=False)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
mode = Column(String, nullable=False) # "A" or "B"
|
||||
prompt = Column(Text, nullable=False)
|
||||
selection_type = Column(String, nullable=False) # "rectangle", "ellipse", "lasso"
|
||||
bbox_json = Column(Text, nullable=False) # JSON string of {x, y, width, height}
|
||||
feather_px = Column(Integer, default=0)
|
||||
ai_provider = Column(String, nullable=False)
|
||||
status = Column(String, nullable=False) # "pending", "processing", "completed", "failed"
|
||||
error_message = Column(Text, nullable=True)
|
||||
|
||||
# Relationships
|
||||
project = relationship("Project", back_populates="edits")
|
||||
@@ -0,0 +1,37 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Text, Boolean
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime
|
||||
from app.database import Base
|
||||
|
||||
|
||||
class Patch(Base):
|
||||
__tablename__ = "patches"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=True)
|
||||
name = Column(String, nullable=False)
|
||||
description = Column(Text, nullable=True)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
|
||||
# Source information
|
||||
source_type = Column(String, nullable=False) # "ai_generated", "manual_selection", "imported"
|
||||
source_project_id = Column(Integer, ForeignKey("projects.id"), nullable=True)
|
||||
source_edit_id = Column(Integer, ForeignKey("edits.id"), nullable=True)
|
||||
|
||||
# Patch metadata
|
||||
width = Column(Integer, nullable=False)
|
||||
height = Column(Integer, nullable=False)
|
||||
tags = Column(Text, nullable=True) # Comma-separated tags
|
||||
category = Column(String, nullable=True) # "hand", "face", "body", "object", "texture", etc.
|
||||
|
||||
# Is this patch shared/public?
|
||||
is_public = Column(Boolean, default=False)
|
||||
|
||||
# File path (relative to data dir)
|
||||
file_path = Column(String, nullable=False)
|
||||
thumbnail_path = Column(String, nullable=True)
|
||||
|
||||
# Relationships
|
||||
user = relationship("User", back_populates="patches")
|
||||
source_project = relationship("Project")
|
||||
source_edit = relationship("Edit")
|
||||
@@ -0,0 +1,18 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime
|
||||
from app.database import Base
|
||||
|
||||
|
||||
class Project(Base):
|
||||
__tablename__ = "projects"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||||
name = Column(String, nullable=False)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
||||
|
||||
# Relationships
|
||||
user = relationship("User", back_populates="projects")
|
||||
edits = relationship("Edit", back_populates="project", cascade="all, delete-orphan")
|
||||
@@ -0,0 +1,17 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime
|
||||
from app.database import Base
|
||||
|
||||
|
||||
class User(Base):
|
||||
__tablename__ = "users"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
email = Column(String, unique=True, index=True, nullable=False)
|
||||
password_hash = Column(String, nullable=False)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
|
||||
# Relationships
|
||||
projects = relationship("Project", back_populates="user", cascade="all, delete-orphan")
|
||||
patches = relationship("Patch", back_populates="user", cascade="all, delete-orphan")
|
||||
@@ -0,0 +1,989 @@
|
||||
"""
|
||||
AI tools router — LaMa inpaint, background removal, remote generation, config.
|
||||
All endpoints are under /api prefix.
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional
|
||||
import base64
|
||||
import asyncio
|
||||
import json
|
||||
from io import BytesIO
|
||||
|
||||
from app.services.local_inpaint import (
|
||||
lama_inpaint, opencv_inpaint, lama_available, gpu_available, rembg_available,
|
||||
)
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["ai-tools"])
|
||||
|
||||
|
||||
# ─── Request / response models ───────────────────────────────────────────────
|
||||
|
||||
class EraseRequest(BaseModel):
|
||||
image: str # base64
|
||||
mask: str # base64
|
||||
|
||||
|
||||
class InpaintRemoteRequest(BaseModel):
|
||||
image: str
|
||||
mask: str
|
||||
prompt: str
|
||||
negative_prompt: Optional[str] = ""
|
||||
steps: Optional[int] = 30
|
||||
cfg_scale: Optional[float] = 7.5
|
||||
model: Optional[str] = None
|
||||
|
||||
|
||||
class Txt2ImgRequest(BaseModel):
|
||||
prompt: str
|
||||
width: Optional[int] = 1024
|
||||
height: Optional[int] = 1024
|
||||
negative_prompt: Optional[str] = ""
|
||||
steps: Optional[int] = 30
|
||||
cfg_scale: Optional[float] = 7.5
|
||||
model: Optional[str] = None
|
||||
seed: Optional[int] = 0
|
||||
|
||||
|
||||
class Img2ImgRequest(BaseModel):
|
||||
image: str
|
||||
prompt: str
|
||||
strength: Optional[float] = 0.75
|
||||
negative_prompt: Optional[str] = ""
|
||||
steps: Optional[int] = 30
|
||||
cfg_scale: Optional[float] = 7.5
|
||||
model: Optional[str] = None
|
||||
|
||||
|
||||
class OutpaintRequest(BaseModel):
|
||||
image: str
|
||||
direction: str # left | right | top | bottom
|
||||
size: Optional[int] = 256
|
||||
prompt: Optional[str] = ""
|
||||
|
||||
|
||||
class BgRemoveRequest(BaseModel):
|
||||
image: str
|
||||
|
||||
|
||||
# ─── Helpers ─────────────────────────────────────────────────────────────────
|
||||
|
||||
def _decode(b64: str) -> bytes:
|
||||
return base64.b64decode(b64)
|
||||
|
||||
|
||||
def _encode(data: bytes) -> str:
|
||||
return base64.b64encode(data).decode()
|
||||
|
||||
|
||||
def _require_remote(operation: str = None):
|
||||
from app.services.remote_provider import get_remote_provider
|
||||
from app.config import settings
|
||||
provider = get_remote_provider(operation)
|
||||
if provider is None:
|
||||
if (settings.ai_provider or "").lower() == "local_gpu":
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail=(
|
||||
"local_gpu provider failed to load — diffusers may be incompatible with "
|
||||
"the installed PyTorch version. Check container logs for details. "
|
||||
"If you see 'torch has no attribute xpu', rebuild the container from the "
|
||||
"correct branch so the pinned diffusers<0.29.0 is installed."
|
||||
)
|
||||
)
|
||||
op_hint = f"AI_PROVIDER_{operation.upper()} or " if operation else ""
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail=f"No remote AI provider configured for '{operation or 'default'}'. "
|
||||
f"Set {op_hint}AI_PROVIDER in .env (openai / invokeai / comfyui)."
|
||||
)
|
||||
return provider
|
||||
|
||||
|
||||
# ─── Local inpaint endpoints ─────────────────────────────────────────────────
|
||||
|
||||
@router.post("/erase")
|
||||
async def erase(req: EraseRequest):
|
||||
"""
|
||||
Magic eraser: remove object / fill region using LaMa (local, no API key needed).
|
||||
Falls back to OpenCV if LaMa not installed.
|
||||
"""
|
||||
try:
|
||||
image_bytes = _decode(req.image)
|
||||
mask_bytes = _decode(req.mask)
|
||||
|
||||
if lama_available():
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, lama_inpaint, image_bytes, mask_bytes
|
||||
)
|
||||
method = "lama"
|
||||
else:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, opencv_inpaint, image_bytes, mask_bytes
|
||||
)
|
||||
method = "opencv"
|
||||
|
||||
return {"result": _encode(result), "method": method}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/inpaint/lama")
|
||||
async def inpaint_lama(req: EraseRequest):
|
||||
"""LaMa structural inpainting."""
|
||||
if not lama_available():
|
||||
raise HTTPException(status_code=503, detail="simple-lama-inpainting not installed.")
|
||||
try:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, lama_inpaint, _decode(req.image), _decode(req.mask)
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/inpaint/fast")
|
||||
async def inpaint_fast(req: EraseRequest):
|
||||
"""OpenCV fast inpainting (CPU, milliseconds)."""
|
||||
try:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, opencv_inpaint, _decode(req.image), _decode(req.mask)
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/background/remove")
|
||||
async def background_remove(req: BgRemoveRequest):
|
||||
"""Remove background — rembg if available, else U2Net."""
|
||||
try:
|
||||
image_bytes = _decode(req.image)
|
||||
|
||||
# Try rembg first
|
||||
if rembg_available():
|
||||
from app.services.local_inpaint import remove_background_rembg
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, remove_background_rembg, image_bytes
|
||||
)
|
||||
return {"result": _encode(result), "method": "rembg"}
|
||||
|
||||
# Fall back to U2Net (existing implementation)
|
||||
from PIL import Image
|
||||
from io import BytesIO as _BytesIO
|
||||
img = Image.open(_BytesIO(image_bytes)).convert("RGB")
|
||||
from app.routers.tools import _remove_background_u2net
|
||||
result = await _remove_background_u2net(img)
|
||||
return {"result": _encode(result), "method": "u2net"}
|
||||
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# ─── Remote provider endpoints ───────────────────────────────────────────────
|
||||
|
||||
@router.post("/inpaint/remote")
|
||||
async def inpaint_remote(req: InpaintRemoteRequest):
|
||||
"""Inpaint via configured remote provider (InvokeAI / ComfyUI / OpenAI)."""
|
||||
provider = _require_remote("inpaint")
|
||||
try:
|
||||
params = {
|
||||
"negative_prompt": req.negative_prompt or "",
|
||||
"steps": req.steps,
|
||||
"cfg_scale": req.cfg_scale,
|
||||
}
|
||||
if req.model:
|
||||
params["model"] = req.model
|
||||
result = await provider.inpaint(_decode(req.image), _decode(req.mask), req.prompt, params)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/generate/progress")
|
||||
async def generation_progress_stream():
|
||||
"""
|
||||
SSE stream of local GPU pipeline inference progress.
|
||||
Events are JSON arrays of pipeline state objects, emitted every 200 ms.
|
||||
Each object: {pipeline, state, step, total_steps, progress, message, model_id, …}
|
||||
Clients open this with EventSource before firing a generation POST,
|
||||
then close it when the POST resolves.
|
||||
"""
|
||||
from app.services.local_diffusion import get_all_model_states
|
||||
|
||||
async def event_gen():
|
||||
try:
|
||||
while True:
|
||||
states = get_all_model_states()
|
||||
yield f"data: {json.dumps(states)}\n\n"
|
||||
await asyncio.sleep(0.2)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
return StreamingResponse(
|
||||
event_gen(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"X-Accel-Buffering": "no",
|
||||
"Connection": "keep-alive",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/generate/txt2img")
|
||||
async def txt2img(req: Txt2ImgRequest):
|
||||
"""Text-to-image via configured remote provider."""
|
||||
provider = _require_remote("txt2img")
|
||||
try:
|
||||
params = {
|
||||
"negative_prompt": req.negative_prompt or "",
|
||||
"steps": req.steps,
|
||||
"cfg_scale": req.cfg_scale,
|
||||
"seed": req.seed or 0,
|
||||
}
|
||||
if req.model:
|
||||
params["model"] = req.model
|
||||
result = await provider.txt2img(req.prompt, req.width, req.height, params)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/generate/img2img")
|
||||
async def img2img(req: Img2ImgRequest):
|
||||
"""Image-to-image via configured remote provider."""
|
||||
provider = _require_remote("img2img")
|
||||
try:
|
||||
params = {
|
||||
"negative_prompt": req.negative_prompt or "",
|
||||
"steps": req.steps,
|
||||
"cfg_scale": req.cfg_scale,
|
||||
}
|
||||
if req.model:
|
||||
params["model"] = req.model
|
||||
result = await provider.img2img(_decode(req.image), req.prompt, req.strength, params)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/generate/outpaint")
|
||||
async def outpaint(req: OutpaintRequest):
|
||||
"""Expand canvas in given direction via remote provider."""
|
||||
provider = _require_remote("outpaint")
|
||||
if req.direction not in ("left", "right", "top", "bottom"):
|
||||
raise HTTPException(status_code=400, detail="direction must be left/right/top/bottom")
|
||||
try:
|
||||
result = await provider.outpaint(_decode(req.image), req.direction, req.size, req.prompt or "")
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# ─── Config / capabilities ────────────────────────────────────────────────────
|
||||
|
||||
class ConfigUpdateRequest(BaseModel):
|
||||
ai_provider: Optional[str] = None
|
||||
# Per-operation overrides (blank = use default)
|
||||
ai_provider_inpaint: Optional[str] = None
|
||||
ai_provider_txt2img: Optional[str] = None
|
||||
ai_provider_img2img: Optional[str] = None
|
||||
ai_provider_outpaint: Optional[str] = None
|
||||
# Credentials / URLs
|
||||
openai_api_key: Optional[str] = None
|
||||
openai_model: Optional[str] = None
|
||||
invokeai_url: Optional[str] = None
|
||||
invokeai_default_model: Optional[str] = None
|
||||
comfyui_url: Optional[str] = None
|
||||
comfyui_default_model: Optional[str] = None
|
||||
replicate_api_key: Optional[str] = None
|
||||
stability_api_key: Optional[str] = None
|
||||
|
||||
|
||||
@router.post("/config")
|
||||
async def update_config(req: ConfigUpdateRequest):
|
||||
"""
|
||||
Apply runtime provider settings (no restart needed).
|
||||
Values are applied to the live settings object in-process.
|
||||
They do NOT persist across restarts — set them in .env for permanence.
|
||||
"""
|
||||
from app.config import settings
|
||||
|
||||
_str_fields = [
|
||||
"ai_provider", "ai_provider_inpaint", "ai_provider_txt2img",
|
||||
"ai_provider_img2img", "ai_provider_outpaint",
|
||||
"openai_api_key", "openai_model",
|
||||
"invokeai_url", "invokeai_default_model",
|
||||
"comfyui_url", "comfyui_default_model",
|
||||
"replicate_api_key", "stability_api_key",
|
||||
]
|
||||
for field in _str_fields:
|
||||
val = getattr(req, field, None)
|
||||
if val is not None:
|
||||
setattr(settings, field, val)
|
||||
|
||||
return {
|
||||
"status": "ok",
|
||||
"ai_provider": settings.ai_provider,
|
||||
"overrides": {
|
||||
"inpaint": settings.ai_provider_inpaint or None,
|
||||
"txt2img": settings.ai_provider_txt2img or None,
|
||||
"img2img": settings.ai_provider_img2img or None,
|
||||
"outpaint": settings.ai_provider_outpaint or None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
async def _check_provider(operation: str) -> dict:
|
||||
"""Health-check the provider for a specific operation."""
|
||||
from app.services.remote_provider import get_remote_provider
|
||||
try:
|
||||
p = get_remote_provider(operation)
|
||||
if p is None:
|
||||
return {"provider": None, "healthy": False}
|
||||
healthy = await asyncio.wait_for(p.health(), timeout=5.0)
|
||||
return {"provider": p.__class__.__name__.replace("Provider", "").lower(), "healthy": healthy}
|
||||
except Exception:
|
||||
return {"provider": None, "healthy": False}
|
||||
|
||||
|
||||
@router.get("/config")
|
||||
async def get_config():
|
||||
"""
|
||||
Return capability flags so the frontend can show/hide tools.
|
||||
Includes per-operation provider assignments and health status.
|
||||
"""
|
||||
from app.config import settings
|
||||
|
||||
# Run health checks for each operation concurrently
|
||||
ops = ["inpaint", "txt2img", "img2img", "outpaint"]
|
||||
results = await asyncio.gather(*[_check_provider(op) for op in ops])
|
||||
op_status = dict(zip(ops, results))
|
||||
|
||||
# Default provider for display (used when no per-op override)
|
||||
default_name = (settings.ai_provider or "").lower() or None
|
||||
|
||||
from app.services.gpu_detect import get_cached_gpu_info
|
||||
gpu_info = get_cached_gpu_info()
|
||||
|
||||
return {
|
||||
"local": {
|
||||
"lama": lama_available(),
|
||||
"rembg": rembg_available(),
|
||||
"opencv": True,
|
||||
"gpu_detected": gpu_available(),
|
||||
"gpu_backend": gpu_info.backend,
|
||||
"gpu_device": gpu_info.device_name,
|
||||
"gpu_vram_total": gpu_info.vram_total_gb,
|
||||
"gpu_vram_free": gpu_info.vram_free_gb,
|
||||
"gpu_cc": gpu_info.compute_capability,
|
||||
"gpu_fp16": gpu_info.fp16,
|
||||
"gpu_bf16": gpu_info.bf16,
|
||||
"gpu_fp8": gpu_info.fp8,
|
||||
"gpu_tensor_cores": gpu_info.tensor_cores,
|
||||
"gpu_tier": gpu_info.tier,
|
||||
"gpu_eff_vram": gpu_info.effective_vram_gb,
|
||||
"local_gpu_available": gpu_info.backend in ("cuda", "mps"),
|
||||
"local_gpu_capabilities": gpu_info.capabilities,
|
||||
"local_gpu_warnings": gpu_info.warnings,
|
||||
},
|
||||
"remote": {
|
||||
"default_provider": default_name,
|
||||
# Legacy field kept for backwards compat with badge/capabilities checks
|
||||
"provider": default_name,
|
||||
"healthy": any(v["healthy"] for v in op_status.values()),
|
||||
"operations": op_status,
|
||||
"overrides": {
|
||||
"inpaint": settings.ai_provider_inpaint or None,
|
||||
"txt2img": settings.ai_provider_txt2img or None,
|
||||
"img2img": settings.ai_provider_img2img or None,
|
||||
"outpaint": settings.ai_provider_outpaint or None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
# ─── Selection image operations ─────────────────────────────────────────────
|
||||
|
||||
class ScaleSelectionRequest(BaseModel):
|
||||
image: str # base64 full canvas
|
||||
mask: str # base64 selection mask (white = object)
|
||||
scale_pct: float = 103.0 # 103 = 3% bigger, 95 = 5% smaller
|
||||
|
||||
|
||||
class AiEditRegionRequest(BaseModel):
|
||||
image: str
|
||||
mask: str
|
||||
instruction: str
|
||||
negative_prompt: str = ""
|
||||
steps: int = 30
|
||||
cfg_scale: float = 7.5
|
||||
|
||||
|
||||
class PasteIntoSelectionRequest(BaseModel):
|
||||
image: str # base64 target canvas
|
||||
mask: str # base64 selection mask
|
||||
paste_image: str # base64 image to paste
|
||||
|
||||
|
||||
@router.post("/image/scale-selection")
|
||||
async def scale_selection(req: ScaleSelectionRequest):
|
||||
"""
|
||||
Scale the object selected by mask by scale_pct%, AI-fill the exposed gap.
|
||||
Works purely with local tools (LaMa/OpenCV) — no remote provider needed.
|
||||
"""
|
||||
try:
|
||||
import numpy as np
|
||||
from PIL import Image, ImageFilter
|
||||
except ImportError:
|
||||
raise HTTPException(status_code=500, detail="PIL/numpy not available")
|
||||
|
||||
img = Image.open(BytesIO(_decode(req.image))).convert("RGB")
|
||||
mask = Image.open(BytesIO(_decode(req.mask))).convert("L")
|
||||
if img.size != mask.size:
|
||||
mask = mask.resize(img.size, Image.LANCZOS)
|
||||
|
||||
mask_arr = np.array(mask)
|
||||
ys, xs = np.where(mask_arr > 128)
|
||||
if len(xs) == 0:
|
||||
raise HTTPException(status_code=400, detail="Empty mask — nothing to scale")
|
||||
|
||||
minx, maxx = int(xs.min()), int(xs.max())
|
||||
miny, maxy = int(ys.min()), int(ys.max())
|
||||
cx, cy = (minx + maxx) / 2.0, (miny + maxy) / 2.0
|
||||
obj_w, obj_h = maxx - minx + 1, maxy - miny + 1
|
||||
|
||||
scale = req.scale_pct / 100.0
|
||||
new_w = max(1, round(obj_w * scale))
|
||||
new_h = max(1, round(obj_h * scale))
|
||||
|
||||
# Extract masked object crop (RGBA with mask as alpha)
|
||||
img_rgba = img.convert("RGBA")
|
||||
obj_crop = img_rgba.crop((minx, miny, maxx + 1, maxy + 1))
|
||||
mask_crop = mask.crop((minx, miny, maxx + 1, maxy + 1))
|
||||
r, g, b, _ = obj_crop.split()
|
||||
obj_masked = Image.merge("RGBA", (r, g, b, mask_crop))
|
||||
scaled_obj = obj_masked.resize((new_w, new_h), Image.LANCZOS)
|
||||
|
||||
# AI-fill the original mask area (gap) with LaMa/OpenCV
|
||||
gap_mask = mask.filter(ImageFilter.MaxFilter(9)) # expand ~4px for clean seam
|
||||
gap_bytes = BytesIO()
|
||||
img.save(gap_bytes, format="PNG")
|
||||
gap_mask_bytes = BytesIO()
|
||||
gap_mask.save(gap_mask_bytes, format="PNG")
|
||||
|
||||
try:
|
||||
if lama_available():
|
||||
filled_bytes = await asyncio.get_event_loop().run_in_executor(
|
||||
None, lama_inpaint, gap_bytes.getvalue(), gap_mask_bytes.getvalue()
|
||||
)
|
||||
else:
|
||||
filled_bytes = await asyncio.get_event_loop().run_in_executor(
|
||||
None, opencv_inpaint, gap_bytes.getvalue(), gap_mask_bytes.getvalue()
|
||||
)
|
||||
filled = Image.open(BytesIO(filled_bytes)).convert("RGBA")
|
||||
except Exception as exc:
|
||||
print(f"[scale-selection] fill fallback: {exc}")
|
||||
filled = img.convert("RGBA")
|
||||
|
||||
# Paste scaled object centered on original centroid
|
||||
px = round(cx - new_w / 2)
|
||||
py = round(cy - new_h / 2)
|
||||
result = filled.copy()
|
||||
result.paste(scaled_obj, (px, py), scaled_obj.split()[3])
|
||||
|
||||
out = BytesIO()
|
||||
result.convert("RGB").save(out, format="PNG")
|
||||
return {"result": _encode(out.getvalue())}
|
||||
|
||||
|
||||
@router.post("/image/ai-edit-region")
|
||||
async def ai_edit_region(req: AiEditRegionRequest):
|
||||
"""
|
||||
AI-edit the selected region using the configured inpaint provider.
|
||||
Works with local_gpu, InvokeAI, ComfyUI, or OpenAI.
|
||||
"""
|
||||
provider = _require_remote("inpaint")
|
||||
try:
|
||||
result_bytes = await provider.inpaint(
|
||||
_decode(req.image),
|
||||
_decode(req.mask),
|
||||
req.instruction,
|
||||
{"negative_prompt": req.negative_prompt, "steps": req.steps, "cfg_scale": req.cfg_scale},
|
||||
)
|
||||
except Exception as exc:
|
||||
import traceback; traceback.print_exc()
|
||||
msg = str(exc)
|
||||
if "Errno -3" in msg or "Name or service not known" in msg or "ConnectError" in msg:
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail=(
|
||||
"AI model files not yet downloaded — container DNS appears to be blocked. "
|
||||
"Fix: sudo iptables -I DOCKER-USER -p udp --dport 53 -j ACCEPT on the host, "
|
||||
"or pre-download the model: pip install huggingface-hub && "
|
||||
"huggingface-cli download diffusers/stable-diffusion-xl-1.0-inpainting-0.1 "
|
||||
"--cache-dir ./data/hf_cache"
|
||||
)
|
||||
)
|
||||
raise HTTPException(status_code=500, detail=msg)
|
||||
return {"result": _encode(result_bytes)}
|
||||
|
||||
|
||||
@router.post("/image/paste-into-selection")
|
||||
async def paste_into_selection(req: PasteIntoSelectionRequest):
|
||||
"""
|
||||
Scale a clipboard image to the selection bounding box, mask it to the
|
||||
selection shape, and composite it over the original canvas.
|
||||
"""
|
||||
try:
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
except ImportError:
|
||||
raise HTTPException(status_code=500, detail="PIL/numpy not available")
|
||||
|
||||
img = Image.open(BytesIO(_decode(req.image))).convert("RGBA")
|
||||
mask = Image.open(BytesIO(_decode(req.mask))).convert("L")
|
||||
paste_img = Image.open(BytesIO(_decode(req.paste_image))).convert("RGBA")
|
||||
|
||||
if img.size != mask.size:
|
||||
mask = mask.resize(img.size, Image.LANCZOS)
|
||||
|
||||
mask_arr = np.array(mask)
|
||||
ys, xs = np.where(mask_arr > 128)
|
||||
if len(xs) == 0:
|
||||
raise HTTPException(status_code=400, detail="Empty mask")
|
||||
|
||||
minx, maxx = int(xs.min()), int(xs.max())
|
||||
miny, maxy = int(ys.min()), int(ys.max())
|
||||
target_w = maxx - minx + 1
|
||||
target_h = maxy - miny + 1
|
||||
|
||||
# Scale clipboard image to fit the selection bounding box
|
||||
paste_scaled = paste_img.resize((target_w, target_h), Image.LANCZOS)
|
||||
|
||||
# Clip paste to selection shape using mask
|
||||
mask_crop = mask.crop((minx, miny, maxx + 1, maxy + 1))
|
||||
r, g, b, a = paste_scaled.split()
|
||||
mask_np = np.array(mask_crop)
|
||||
alpha_np = np.array(a)
|
||||
combined = (alpha_np.astype(np.uint16) * mask_np.astype(np.uint16) // 255).astype(np.uint8)
|
||||
paste_final = Image.merge("RGBA", (r, g, b, Image.fromarray(combined)))
|
||||
|
||||
result = img.copy()
|
||||
result.paste(paste_final, (minx, miny), paste_final.split()[3])
|
||||
|
||||
out = BytesIO()
|
||||
result.convert("RGB").save(out, format="PNG")
|
||||
return {"result": _encode(out.getvalue())}
|
||||
|
||||
|
||||
# ─── SAM (Segment Anything) ──────────────────────────────────────────────────
|
||||
|
||||
class SegmentPointRequest(BaseModel):
|
||||
image: str # base64 PNG/JPEG
|
||||
points: list[list[int]] # [[x, y], ...] original image coords
|
||||
labels: list[int] # 1=include, 0=exclude — same length as points
|
||||
|
||||
|
||||
@router.post("/segment/point")
|
||||
async def segment_point(req: SegmentPointRequest):
|
||||
"""
|
||||
Run SAM point-prompt segmentation.
|
||||
Returns a binary mask PNG (white = selected area).
|
||||
Auto-downloads the SAM ViT-B model (~375 MB) on first call.
|
||||
"""
|
||||
if not req.points:
|
||||
raise HTTPException(status_code=400, detail="At least one point required.")
|
||||
if len(req.points) != len(req.labels):
|
||||
raise HTTPException(status_code=400, detail="points and labels must have the same length.")
|
||||
|
||||
try:
|
||||
image_bytes = base64.b64decode(req.image)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Could not decode image: {e}")
|
||||
|
||||
from app.services.sam_service import predict_points, get_install_status
|
||||
try:
|
||||
mask_bytes = await predict_points(
|
||||
image_bytes,
|
||||
[tuple(p) for p in req.points],
|
||||
req.labels,
|
||||
)
|
||||
return {
|
||||
"mask": base64.b64encode(mask_bytes).decode(),
|
||||
"sam_install": get_install_status(),
|
||||
}
|
||||
except RuntimeError as e:
|
||||
raise HTTPException(status_code=503, detail=str(e))
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/segment/install-status")
|
||||
def segment_install_status():
|
||||
"""Poll SAM model download progress."""
|
||||
from app.services.sam_service import get_install_status, sam_model_available
|
||||
status = get_install_status()
|
||||
status["model_ready"] = sam_model_available()
|
||||
return status
|
||||
|
||||
|
||||
@router.post("/segment/install")
|
||||
async def segment_install():
|
||||
"""Trigger SAM model download explicitly (also auto-triggered on first /segment/point call)."""
|
||||
from app.services.sam_service import ensure_sam_installed, get_install_status
|
||||
asyncio.create_task(ensure_sam_installed())
|
||||
return get_install_status()
|
||||
|
||||
|
||||
# ─── Enhance ─────────────────────────────────────────────────────────────────
|
||||
|
||||
import io as _io
|
||||
import numpy as _np
|
||||
import cv2 as _cv2
|
||||
from PIL import Image as _Image
|
||||
|
||||
class EnhanceRequest(BaseModel):
|
||||
image: str # base64
|
||||
strength: float = 1.0
|
||||
|
||||
|
||||
def _enhance_image(image_bytes: bytes, strength: float) -> bytes:
|
||||
"""
|
||||
Apply a chain of non-AI image enhancements, each blended with `strength` (0–1).
|
||||
|
||||
Steps:
|
||||
1. Auto white balance (gray-world)
|
||||
2. CLAHE on L channel of LAB colorspace
|
||||
3. Auto saturation boost in HSV (×1.15, clamped)
|
||||
4. Mild unsharp mask (gaussian sigma=1.0, delta weight=0.3)
|
||||
"""
|
||||
strength = max(0.0, min(1.0, float(strength)))
|
||||
|
||||
# Decode to RGB numpy array
|
||||
pil = _Image.open(_io.BytesIO(image_bytes)).convert("RGB")
|
||||
orig = _np.array(pil, dtype=_np.float32) # H×W×3, float [0,255]
|
||||
|
||||
img = orig.copy()
|
||||
|
||||
# ── Step 1: Auto white balance (gray-world) ──────────────────────────────
|
||||
mean_r = img[:, :, 0].mean()
|
||||
mean_g = img[:, :, 1].mean()
|
||||
mean_b = img[:, :, 2].mean()
|
||||
overall_mean = (mean_r + mean_g + mean_b) / 3.0
|
||||
|
||||
def _scale(channel, channel_mean):
|
||||
if channel_mean == 0:
|
||||
return channel
|
||||
return channel * (overall_mean / channel_mean)
|
||||
|
||||
wb = img.copy()
|
||||
wb[:, :, 0] = _np.clip(_scale(img[:, :, 0], mean_r), 0, 255)
|
||||
wb[:, :, 1] = _np.clip(_scale(img[:, :, 1], mean_g), 0, 255)
|
||||
wb[:, :, 2] = _np.clip(_scale(img[:, :, 2], mean_b), 0, 255)
|
||||
|
||||
img = (orig + strength * (wb - orig)).clip(0, 255)
|
||||
|
||||
# ── Step 2: CLAHE on L channel (LAB) ────────────────────────────────────
|
||||
img_u8 = img.astype(_np.uint8)
|
||||
lab = _cv2.cvtColor(img_u8, _cv2.COLOR_RGB2LAB)
|
||||
clahe = _cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||||
l_orig = lab[:, :, 0].copy()
|
||||
lab[:, :, 0] = clahe.apply(l_orig)
|
||||
# Blend L channel back using strength
|
||||
lab_blended = lab.copy()
|
||||
lab_blended[:, :, 0] = (l_orig + strength * (lab[:, :, 0].astype(_np.float32) - l_orig.astype(_np.float32))).clip(0, 255).astype(_np.uint8)
|
||||
img = _cv2.cvtColor(lab_blended, _cv2.COLOR_LAB2RGB).astype(_np.float32)
|
||||
|
||||
# ── Step 3: Auto saturation boost (HSV, ×1.15) ──────────────────────────
|
||||
img_u8 = img.astype(_np.uint8)
|
||||
hsv = _cv2.cvtColor(img_u8, _cv2.COLOR_RGB2HSV).astype(_np.float32)
|
||||
s_orig = hsv[:, :, 1].copy()
|
||||
s_boosted = _np.clip(s_orig * 1.15, 0, 255)
|
||||
hsv[:, :, 1] = s_orig + strength * (s_boosted - s_orig)
|
||||
hsv = hsv.clip(0, 255).astype(_np.uint8)
|
||||
img = _cv2.cvtColor(hsv, _cv2.COLOR_HSV2RGB).astype(_np.float32)
|
||||
|
||||
# ── Step 4: Mild unsharp mask (sigma=1.0, delta weight=0.3) ─────────────
|
||||
img_u8 = img.astype(_np.uint8)
|
||||
blurred = _cv2.GaussianBlur(img_u8, (0, 0), sigmaX=1.0)
|
||||
sharpness_delta = img_u8.astype(_np.float32) - blurred.astype(_np.float32)
|
||||
sharpened = img_u8.astype(_np.float32) + 0.3 * sharpness_delta * strength
|
||||
img = sharpened.clip(0, 255)
|
||||
|
||||
# Encode result as PNG
|
||||
result_pil = _Image.fromarray(img.astype(_np.uint8), mode="RGB")
|
||||
buf = _io.BytesIO()
|
||||
result_pil.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
@router.post("/enhance")
|
||||
async def enhance(req: EnhanceRequest):
|
||||
"""
|
||||
Non-AI image enhancement: auto white balance, CLAHE, saturation boost,
|
||||
and unsharp mask. Each step is blended proportionally to `strength` (0–1).
|
||||
"""
|
||||
try:
|
||||
image_bytes = _decode(req.image)
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, _enhance_image, image_bytes, req.strength
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# ─── Subject replace ─────────────────────────────────────────────────────────
|
||||
|
||||
class ExtractSubjectRequest(BaseModel):
|
||||
image: str # base64
|
||||
|
||||
|
||||
class ReplaceSubjectRequest(BaseModel):
|
||||
background_image: str # base64 — image whose background we keep
|
||||
subject_image: str # base64 — image whose subject we extract
|
||||
mask: Optional[str] = None # base64 — white = where the subject should land
|
||||
match_colors: bool = True # blend subject color stats toward background
|
||||
|
||||
|
||||
def _extract_subject_bytes(image_bytes: bytes) -> bytes:
|
||||
"""Remove background from image using rembg; return RGBA PNG bytes."""
|
||||
if rembg_available():
|
||||
return remove_background_rembg(image_bytes)
|
||||
raise RuntimeError(
|
||||
"rembg is not installed. Run: pip install rembg (or add it to requirements.txt)"
|
||||
)
|
||||
|
||||
|
||||
def _color_transfer_lab(subj_rgba: "Image", bg_rgb: "Image", blend: float = 0.45) -> "Image":
|
||||
"""
|
||||
Partial LAB color transfer: nudge subject color statistics 'blend' fraction
|
||||
toward the background's statistics so it looks like it belongs in the scene.
|
||||
"""
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
src_arr = np.array(subj_rgba.convert("RGB"), dtype=np.float32)
|
||||
tgt_arr = np.array(bg_rgb.convert("RGB"), dtype=np.float32)
|
||||
|
||||
alpha = np.array(subj_rgba.split()[3])
|
||||
subject_mask = alpha > 10
|
||||
|
||||
if not subject_mask.any():
|
||||
return subj_rgba
|
||||
|
||||
src_lab = cv2.cvtColor(src_arr.astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
|
||||
tgt_lab = cv2.cvtColor(tgt_arr.astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
|
||||
|
||||
for ch in range(3):
|
||||
src_ch = src_lab[:, :, ch]
|
||||
src_pixels = src_ch[subject_mask]
|
||||
tgt_pixels = tgt_lab[:, :, ch].flatten()
|
||||
|
||||
src_mean, src_std = float(src_pixels.mean()), float(src_pixels.std()) + 1e-6
|
||||
tgt_mean, tgt_std = float(tgt_pixels.mean()), float(tgt_pixels.std()) + 1e-6
|
||||
|
||||
adjusted_std = src_std + blend * (tgt_std - src_std)
|
||||
adjusted = (src_ch - src_mean) * (adjusted_std / src_std) + src_mean + blend * (tgt_mean - src_mean)
|
||||
src_lab[:, :, ch] = np.clip(adjusted, 0, 255)
|
||||
|
||||
result_rgb = cv2.cvtColor(src_lab.astype(np.uint8), cv2.COLOR_LAB2RGB)
|
||||
r, g, b = result_rgb[:, :, 0], result_rgb[:, :, 1], result_rgb[:, :, 2]
|
||||
return Image.merge("RGBA", [
|
||||
Image.fromarray(r), Image.fromarray(g),
|
||||
Image.fromarray(b), Image.fromarray(alpha),
|
||||
])
|
||||
|
||||
|
||||
def _do_replace_subject(
|
||||
bg_bytes: bytes,
|
||||
subj_bytes: bytes,
|
||||
mask_bytes: Optional[bytes],
|
||||
match_colors: bool,
|
||||
) -> bytes:
|
||||
"""Core compositing: extract subject → scale → color-match → paste onto background."""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
bg_img = Image.open(BytesIO(bg_bytes)).convert("RGBA")
|
||||
|
||||
subj_rgba = Image.open(BytesIO(_extract_subject_bytes(subj_bytes))).convert("RGBA")
|
||||
|
||||
# Determine target placement bounding box from mask or full canvas
|
||||
if mask_bytes:
|
||||
mask_img = Image.open(BytesIO(mask_bytes)).convert("L")
|
||||
if mask_img.size != bg_img.size:
|
||||
mask_img = mask_img.resize(bg_img.size, Image.LANCZOS)
|
||||
mask_arr = np.array(mask_img)
|
||||
ys, xs = np.where(mask_arr > 128)
|
||||
else:
|
||||
mask_img = None
|
||||
mask_arr = None
|
||||
ys, xs = np.array([]), np.array([])
|
||||
|
||||
if len(xs) > 0:
|
||||
minx, maxx = int(xs.min()), int(xs.max())
|
||||
miny, maxy = int(ys.min()), int(ys.max())
|
||||
else:
|
||||
minx, miny = 0, 0
|
||||
maxx, maxy = bg_img.width - 1, bg_img.height - 1
|
||||
|
||||
target_w = maxx - minx + 1
|
||||
target_h = maxy - miny + 1
|
||||
|
||||
# Scale subject to fit target area, preserving aspect ratio
|
||||
sw, sh = subj_rgba.size
|
||||
scale = min(target_w / sw, target_h / sh)
|
||||
new_w = max(1, round(sw * scale))
|
||||
new_h = max(1, round(sh * scale))
|
||||
subj_scaled = subj_rgba.resize((new_w, new_h), Image.LANCZOS)
|
||||
|
||||
# Optional color transfer to blend lighting/tone
|
||||
if match_colors:
|
||||
subj_scaled = _color_transfer_lab(subj_scaled, bg_img.convert("RGB"))
|
||||
|
||||
# Center in target area
|
||||
px = minx + (target_w - new_w) // 2
|
||||
py = miny + (target_h - new_h) // 2
|
||||
|
||||
result = bg_img.copy()
|
||||
|
||||
if mask_img is not None and len(xs) > 0:
|
||||
# Build a full-canvas RGBA layer for the subject
|
||||
subj_canvas = Image.new("RGBA", bg_img.size, (0, 0, 0, 0))
|
||||
subj_canvas.paste(subj_scaled, (px, py), subj_scaled.split()[3])
|
||||
# Clip subject's alpha to the selection mask
|
||||
sc_arr = np.array(subj_canvas)
|
||||
sc_arr[:, :, 3] = np.minimum(sc_arr[:, :, 3], mask_arr).astype(np.uint8)
|
||||
subj_canvas = Image.fromarray(sc_arr)
|
||||
result.paste(subj_canvas, (0, 0), subj_canvas.split()[3])
|
||||
else:
|
||||
result.paste(subj_scaled, (px, py), subj_scaled.split()[3])
|
||||
|
||||
out = BytesIO()
|
||||
result.convert("RGB").save(out, format="PNG")
|
||||
return out.getvalue()
|
||||
|
||||
|
||||
@router.post("/image/extract-subject")
|
||||
async def extract_subject(req: ExtractSubjectRequest):
|
||||
"""
|
||||
Remove background from an image and return the subject with transparency (RGBA PNG).
|
||||
Uses rembg (AI-powered) when available.
|
||||
"""
|
||||
try:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, _extract_subject_bytes, _decode(req.image)
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/image/replace-subject")
|
||||
async def replace_subject(req: ReplaceSubjectRequest):
|
||||
"""
|
||||
Extract the primary subject from `subject_image` (via rembg background removal),
|
||||
scale it to fit the `mask` selection on `background_image`, apply optional LAB
|
||||
color transfer for lighting consistency, and composite the result.
|
||||
|
||||
Returns the composited image as base64 PNG.
|
||||
"""
|
||||
try:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None,
|
||||
_do_replace_subject,
|
||||
_decode(req.background_image),
|
||||
_decode(req.subject_image),
|
||||
_decode(req.mask) if req.mask else None,
|
||||
req.match_colors,
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# ─── Extract colors ───────────────────────────────────────────────────────────
|
||||
|
||||
class ExtractColorsRequest(BaseModel):
|
||||
image: str # base64
|
||||
count: int = 6
|
||||
|
||||
|
||||
def _extract_colors(image_bytes: bytes, count: int) -> list[str]:
|
||||
"""
|
||||
Resize image to 150×150, k-means cluster pixels into `count` groups
|
||||
using pure numpy (no sklearn dependency), return hex strings by frequency.
|
||||
"""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from io import BytesIO
|
||||
|
||||
count = max(1, min(count, 32))
|
||||
|
||||
pil = Image.open(BytesIO(image_bytes)).convert("RGB").resize((150, 150))
|
||||
pixels = np.array(pil, dtype=np.float32).reshape(-1, 3) # (22500, 3)
|
||||
n = len(pixels)
|
||||
|
||||
# Initialise centers with k-means++ seeding
|
||||
rng = np.random.default_rng(42)
|
||||
centers = [pixels[rng.integers(n)]]
|
||||
for _ in range(count - 1):
|
||||
dists = np.min([np.sum((pixels - c) ** 2, axis=1) for c in centers], axis=0)
|
||||
probs = dists / dists.sum()
|
||||
centers.append(pixels[rng.choice(n, p=probs)])
|
||||
centers = np.array(centers)
|
||||
|
||||
labels = np.zeros(n, dtype=np.int32)
|
||||
for _ in range(20): # max 20 iterations
|
||||
# Assign each pixel to nearest center
|
||||
dists = np.sum((pixels[:, None] - centers[None]) ** 2, axis=2) # (n, k)
|
||||
new_labels = np.argmin(dists, axis=1)
|
||||
if np.all(new_labels == labels):
|
||||
break
|
||||
labels = new_labels
|
||||
# Recompute centers
|
||||
for k in range(count):
|
||||
mask = labels == k
|
||||
if mask.any():
|
||||
centers[k] = pixels[mask].mean(axis=0)
|
||||
|
||||
counts = np.bincount(labels, minlength=count)
|
||||
order = np.argsort(-counts)
|
||||
|
||||
return [
|
||||
"#{:02x}{:02x}{:02x}".format(*centers[i].astype(int).clip(0, 255))
|
||||
for i in order
|
||||
]
|
||||
|
||||
|
||||
@router.post("/extract-colors")
|
||||
async def extract_colors(req: ExtractColorsRequest):
|
||||
"""
|
||||
Extract dominant colors from an image using k-means clustering.
|
||||
Returns hex color strings sorted by frequency (most dominant first).
|
||||
"""
|
||||
try:
|
||||
image_bytes = _decode(req.image)
|
||||
colors = await asyncio.get_event_loop().run_in_executor(
|
||||
None, _extract_colors, image_bytes, req.count
|
||||
)
|
||||
return {"colors": colors}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
@@ -0,0 +1,171 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks
|
||||
from sqlalchemy.orm import Session
|
||||
import json
|
||||
|
||||
from app.database import get_db
|
||||
from app.models.project import Project
|
||||
from app.models.edit import Edit
|
||||
from app.schemas import EditRequest, EditResponse, StatusResponse
|
||||
from app.services.edit_service import EditService
|
||||
from app.config import settings
|
||||
|
||||
router = APIRouter(prefix="/edits", tags=["edits"])
|
||||
|
||||
|
||||
async def process_edit_background(
|
||||
edit_id: int,
|
||||
project_id: int,
|
||||
request: EditRequest,
|
||||
db: Session
|
||||
):
|
||||
"""Background task to process edit"""
|
||||
edit_service = EditService()
|
||||
|
||||
try:
|
||||
# Process the edit
|
||||
result_path = await edit_service.process_edit(
|
||||
project_id=project_id,
|
||||
edit_id=edit_id,
|
||||
prompt=request.prompt,
|
||||
mode=request.mode,
|
||||
selection_type=request.selection_type,
|
||||
bbox=request.bbox,
|
||||
feather_px=request.feather_px,
|
||||
selection_data=request.selection_data
|
||||
)
|
||||
|
||||
# Update edit status
|
||||
edit = db.query(Edit).filter(Edit.id == edit_id).first()
|
||||
if edit:
|
||||
edit.status = "completed"
|
||||
db.commit()
|
||||
|
||||
except Exception as e:
|
||||
# Update edit with error
|
||||
edit = db.query(Edit).filter(Edit.id == edit_id).first()
|
||||
if edit:
|
||||
edit.status = "failed"
|
||||
edit.error_message = str(e)
|
||||
db.commit()
|
||||
|
||||
|
||||
@router.post("/projects/{project_id}/fix", response_model=EditResponse)
|
||||
async def create_edit(
|
||||
project_id: int,
|
||||
request: EditRequest,
|
||||
background_tasks: BackgroundTasks,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Create a new edit request (Fix button)
|
||||
|
||||
This endpoint accepts the selection data and prompt,
|
||||
then processes the edit in the background.
|
||||
"""
|
||||
# Verify project exists
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Validate mode
|
||||
if request.mode not in ["A", "B"]:
|
||||
raise HTTPException(status_code=400, detail="Mode must be 'A' or 'B'")
|
||||
|
||||
# Validate selection type
|
||||
if request.selection_type not in ["rectangle", "ellipse", "lasso"]:
|
||||
raise HTTPException(status_code=400, detail="Invalid selection type")
|
||||
|
||||
# Create edit record
|
||||
edit = Edit(
|
||||
project_id=project_id,
|
||||
mode=request.mode,
|
||||
prompt=request.prompt,
|
||||
selection_type=request.selection_type,
|
||||
bbox_json=json.dumps(request.bbox),
|
||||
feather_px=request.feather_px,
|
||||
ai_provider=settings.ai_provider,
|
||||
status="pending"
|
||||
)
|
||||
db.add(edit)
|
||||
db.commit()
|
||||
db.refresh(edit)
|
||||
|
||||
# Process edit in background
|
||||
background_tasks.add_task(
|
||||
process_edit_background,
|
||||
edit.id,
|
||||
project_id,
|
||||
request,
|
||||
db
|
||||
)
|
||||
|
||||
return edit
|
||||
|
||||
|
||||
@router.get("/{edit_id}", response_model=EditResponse)
|
||||
def get_edit(
|
||||
edit_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get edit details and status"""
|
||||
edit = db.query(Edit).filter(Edit.id == edit_id).first()
|
||||
if not edit:
|
||||
raise HTTPException(status_code=404, detail="Edit not found")
|
||||
return edit
|
||||
|
||||
|
||||
@router.post("/projects/{project_id}/revert/{edit_id}", response_model=StatusResponse)
|
||||
def revert_to_edit(
|
||||
project_id: int,
|
||||
edit_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Revert project to a specific edit"""
|
||||
# Verify project exists
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Verify edit exists and belongs to project
|
||||
edit = db.query(Edit).filter(
|
||||
Edit.id == edit_id,
|
||||
Edit.project_id == project_id
|
||||
).first()
|
||||
if not edit:
|
||||
raise HTTPException(status_code=404, detail="Edit not found")
|
||||
|
||||
# Revert
|
||||
edit_service = EditService()
|
||||
try:
|
||||
result_path = edit_service.revert_to_edit(project_id, edit_id)
|
||||
return StatusResponse(
|
||||
status="success",
|
||||
message=f"Reverted to edit {edit_id}",
|
||||
data={"image_url": f"/projects/{project_id}/current"}
|
||||
)
|
||||
except FileNotFoundError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/projects/{project_id}/reset", response_model=StatusResponse)
|
||||
def reset_to_original(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Reset project to original image"""
|
||||
# Verify project exists
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Reset
|
||||
edit_service = EditService()
|
||||
try:
|
||||
result_path = edit_service.reset_to_original(project_id)
|
||||
return StatusResponse(
|
||||
status="success",
|
||||
message="Reset to original image",
|
||||
data={"image_url": f"/projects/{project_id}/current"}
|
||||
)
|
||||
except FileNotFoundError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
@@ -0,0 +1,176 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, Form
|
||||
from sqlalchemy.orm import Session
|
||||
from typing import Optional
|
||||
from PIL import Image
|
||||
from io import BytesIO
|
||||
import os
|
||||
|
||||
from app.database import get_db
|
||||
from app.models.project import Project
|
||||
from app.schemas import TextToImageRequest, TextToImageResponse
|
||||
from app.services.ai_provider import get_ai_provider
|
||||
from app.services.edit_service import EditService
|
||||
from app.config import settings
|
||||
|
||||
router = APIRouter(prefix="/generate", tags=["generate"])
|
||||
|
||||
|
||||
@router.post("/text-to-image", response_model=TextToImageResponse)
|
||||
async def text_to_image(
|
||||
prompt: str = Form(...),
|
||||
width: int = Form(1024),
|
||||
height: int = Form(1024),
|
||||
negative_prompt: Optional[str] = Form(None),
|
||||
ai_provider: Optional[str] = Form(None),
|
||||
ai_model: Optional[str] = Form(None),
|
||||
create_project: bool = Form(True),
|
||||
project_name: Optional[str] = Form(None),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Generate an image from text prompt
|
||||
|
||||
Args:
|
||||
prompt: Text description of desired image
|
||||
width: Image width (default 1024)
|
||||
height: Image height (default 1024)
|
||||
negative_prompt: What to avoid in generation
|
||||
ai_provider: Override default AI provider
|
||||
ai_model: Specific model to use
|
||||
create_project: Whether to create a new project with the result
|
||||
project_name: Name for the new project (if create_project=True)
|
||||
|
||||
Returns:
|
||||
Generated image info and optionally project details
|
||||
"""
|
||||
|
||||
# Validate dimensions
|
||||
if width < 256 or width > 2048 or height < 256 or height > 2048:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Width and height must be between 256 and 2048"
|
||||
)
|
||||
|
||||
# Get AI provider
|
||||
provider = get_ai_provider(ai_provider, ai_model)
|
||||
|
||||
try:
|
||||
# Generate image
|
||||
image_bytes = await provider.text_to_image(
|
||||
prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
model=ai_model,
|
||||
negative_prompt=negative_prompt
|
||||
)
|
||||
|
||||
project_id = None
|
||||
image_url = None
|
||||
|
||||
if create_project:
|
||||
# Create a new project
|
||||
project = Project(
|
||||
name=project_name or f"Generated: {prompt[:50]}",
|
||||
user_id=None # TODO: Add authentication
|
||||
)
|
||||
db.add(project)
|
||||
db.commit()
|
||||
db.refresh(project)
|
||||
|
||||
project_id = project.id
|
||||
|
||||
# Save image as both original and current
|
||||
edit_service = EditService()
|
||||
edit_service.ensure_project_dir(project_id)
|
||||
|
||||
original_path = edit_service.get_original_image_path(project_id)
|
||||
current_path = edit_service.get_current_image_path(project_id)
|
||||
|
||||
# Save image
|
||||
img = Image.open(BytesIO(image_bytes))
|
||||
img.save(original_path, 'PNG')
|
||||
img.save(current_path, 'PNG')
|
||||
|
||||
image_url = f"/projects/{project_id}/current"
|
||||
|
||||
return TextToImageResponse(
|
||||
status="success",
|
||||
prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
project_id=project_id,
|
||||
image_url=image_url,
|
||||
ai_provider=ai_provider or settings.ai_provider,
|
||||
ai_model=ai_model
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/layer/text-to-image", response_model=TextToImageResponse)
|
||||
async def text_to_image_layer(
|
||||
project_id: int = Form(...),
|
||||
prompt: str = Form(...),
|
||||
width: int = Form(512),
|
||||
height: int = Form(512),
|
||||
x: int = Form(0),
|
||||
y: int = Form(0),
|
||||
negative_prompt: Optional[str] = Form(None),
|
||||
ai_provider: Optional[str] = Form(None),
|
||||
ai_model: Optional[str] = Form(None),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Generate an image as a new layer in an existing project
|
||||
|
||||
This generates a smaller image that can be placed as a layer
|
||||
on top of the current project canvas.
|
||||
"""
|
||||
|
||||
# Verify project exists
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Get AI provider
|
||||
provider = get_ai_provider(ai_provider, ai_model)
|
||||
|
||||
try:
|
||||
# Generate image
|
||||
image_bytes = await provider.text_to_image(
|
||||
prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
model=ai_model,
|
||||
negative_prompt=negative_prompt
|
||||
)
|
||||
|
||||
# Save as temporary layer file
|
||||
edit_service = EditService()
|
||||
layers_dir = edit_service.get_project_dir(project_id) / "layers"
|
||||
layers_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Generate unique layer filename
|
||||
import time
|
||||
layer_filename = f"generated_{int(time.time())}.png"
|
||||
layer_path = layers_dir / layer_filename
|
||||
|
||||
# Save layer image
|
||||
with open(layer_path, 'wb') as f:
|
||||
f.write(image_bytes)
|
||||
|
||||
return TextToImageResponse(
|
||||
status="success",
|
||||
prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
project_id=project_id,
|
||||
image_url=f"/projects/{project_id}/layers/{layer_filename}",
|
||||
layer_position={"x": x, "y": y},
|
||||
ai_provider=ai_provider or settings.ai_provider,
|
||||
ai_model=ai_model
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
@@ -0,0 +1,96 @@
|
||||
"""
|
||||
GPU status and model management endpoints.
|
||||
All under /api/gpu prefix.
|
||||
"""
|
||||
from fastapi import APIRouter
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, List
|
||||
import asyncio
|
||||
|
||||
router = APIRouter(prefix="/api/gpu", tags=["gpu"])
|
||||
|
||||
|
||||
@router.get("/status")
|
||||
async def gpu_status():
|
||||
"""
|
||||
Full GPU capability report: hardware, feature flags, VRAM budget,
|
||||
and which model was selected for each operation.
|
||||
Frontend polls this to show GPU badge and tool availability.
|
||||
"""
|
||||
from app.services.gpu_detect import get_cached_gpu_info
|
||||
from app.services.local_diffusion import get_all_model_states
|
||||
|
||||
info = get_cached_gpu_info()
|
||||
|
||||
return {
|
||||
# Hardware
|
||||
"backend": info.backend,
|
||||
"device_name": info.device_name,
|
||||
"vram_total_gb": info.vram_total_gb,
|
||||
"vram_free_gb": info.vram_free_gb,
|
||||
"compute_capability": info.compute_capability,
|
||||
# Feature flags
|
||||
"fp16": info.fp16,
|
||||
"bf16": info.bf16,
|
||||
"fp8": info.fp8,
|
||||
"int8": info.int8,
|
||||
"tensor_cores": info.tensor_cores,
|
||||
"xformers": info.xformers,
|
||||
# Derived
|
||||
"effective_vram_gb": info.effective_vram_gb,
|
||||
"tier": info.tier,
|
||||
# Selected models per operation
|
||||
"recommended": {
|
||||
op: (
|
||||
{
|
||||
"model_id": spec.model_id,
|
||||
"family": spec.family,
|
||||
"memory_opt": spec.memory_opt,
|
||||
"native_res": spec.native_res,
|
||||
"vram_fp16_gb": spec.vram_fp16_gb,
|
||||
}
|
||||
if spec else None
|
||||
)
|
||||
for op, spec in info.recommended.items()
|
||||
},
|
||||
"pipeline_states": get_all_model_states(),
|
||||
"warnings": info.warnings,
|
||||
"capabilities": info.capabilities,
|
||||
}
|
||||
|
||||
|
||||
class PrefetchRequest(BaseModel):
|
||||
operations: Optional[List[str]] = None
|
||||
|
||||
|
||||
@router.post("/prefetch")
|
||||
async def prefetch_models(req: PrefetchRequest = PrefetchRequest()):
|
||||
"""
|
||||
Eagerly load pipelines into GPU memory for the requested operations.
|
||||
Returns immediately; poll /api/gpu/prefetch-status for progress.
|
||||
Default: inpaint, txt2img, img2img.
|
||||
"""
|
||||
ops = req.operations or ["inpaint", "txt2img", "img2img"]
|
||||
valid = {"inpaint", "txt2img", "img2img", "outpaint", "upscale"}
|
||||
ops = [op for op in ops if op in valid]
|
||||
|
||||
from app.services.local_diffusion import get_local_diffusion_provider
|
||||
provider = get_local_diffusion_provider()
|
||||
|
||||
async def _prefetch():
|
||||
for op in ops:
|
||||
try:
|
||||
await provider._get_pipeline(op)
|
||||
print(f"[gpu] Prefetch complete: {op}")
|
||||
except Exception as exc:
|
||||
print(f"[gpu] Prefetch failed for {op}: {exc}")
|
||||
|
||||
asyncio.create_task(_prefetch())
|
||||
return {"status": "prefetch_started", "operations": ops}
|
||||
|
||||
|
||||
@router.get("/prefetch-status")
|
||||
async def prefetch_status():
|
||||
"""Poll model download / load progress."""
|
||||
from app.services.local_diffusion import get_all_model_states
|
||||
return {"models": get_all_model_states()}
|
||||
@@ -0,0 +1,81 @@
|
||||
from fastapi import APIRouter, HTTPException, Depends
|
||||
from fastapi.responses import FileResponse
|
||||
from sqlalchemy.orm import Session
|
||||
from pathlib import Path
|
||||
|
||||
from app.database import get_db
|
||||
from app.models.project import Project
|
||||
from app.services.edit_service import EditService
|
||||
|
||||
router = APIRouter(prefix="/projects", tags=["images"])
|
||||
|
||||
|
||||
@router.get("/{project_id}/original")
|
||||
def get_original_image(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get the original uploaded image"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
edit_service = EditService()
|
||||
image_path = edit_service.get_original_image_path(project_id)
|
||||
|
||||
if not image_path.exists():
|
||||
raise HTTPException(status_code=404, detail="Original image not found")
|
||||
|
||||
return FileResponse(
|
||||
image_path,
|
||||
media_type="image/png",
|
||||
headers={"Cache-Control": "public, max-age=3600"}
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{project_id}/current")
|
||||
def get_current_image(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get the current edited image"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
edit_service = EditService()
|
||||
image_path = edit_service.get_current_image_path(project_id)
|
||||
|
||||
if not image_path.exists():
|
||||
raise HTTPException(status_code=404, detail="Current image not found")
|
||||
|
||||
return FileResponse(
|
||||
image_path,
|
||||
media_type="image/png",
|
||||
headers={"Cache-Control": "no-cache"}
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{project_id}/history/{edit_id}/result")
|
||||
def get_edit_result(
|
||||
project_id: int,
|
||||
edit_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get the result image from a specific edit"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
edit_service = EditService()
|
||||
edit_dir = edit_service.get_edit_dir(project_id, edit_id)
|
||||
result_path = edit_dir / "result.png"
|
||||
|
||||
if not result_path.exists():
|
||||
raise HTTPException(status_code=404, detail="Edit result not found")
|
||||
|
||||
return FileResponse(
|
||||
result_path,
|
||||
media_type="image/png",
|
||||
headers={"Cache-Control": "public, max-age=3600"}
|
||||
)
|
||||
@@ -0,0 +1,308 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form
|
||||
from fastapi.responses import FileResponse
|
||||
from sqlalchemy.orm import Session
|
||||
from typing import List, Optional
|
||||
import json
|
||||
|
||||
from app.database import get_db
|
||||
from app.models.patch import Patch
|
||||
from app.models.project import Project
|
||||
from app.models.edit import Edit
|
||||
from app.schemas import PatchCreate, PatchResponse, PatchApply, StatusResponse
|
||||
from app.services.patch_library import PatchLibraryService
|
||||
from app.config import settings
|
||||
|
||||
router = APIRouter(prefix="/patches", tags=["patches"])
|
||||
|
||||
|
||||
@router.post("/", response_model=PatchResponse)
|
||||
async def create_patch(
|
||||
name: str = Form(...),
|
||||
description: Optional[str] = Form(None),
|
||||
source_type: str = Form(...),
|
||||
category: Optional[str] = Form(None),
|
||||
tags: Optional[str] = Form(None),
|
||||
source_project_id: Optional[int] = Form(None),
|
||||
source_edit_id: Optional[int] = Form(None),
|
||||
bbox: Optional[str] = Form(None),
|
||||
file: Optional[UploadFile] = File(None),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Create a new patch in the library
|
||||
|
||||
Source types:
|
||||
- ai_generated: From an edit (requires source_edit_id)
|
||||
- manual_selection: Selected from current image (requires source_project_id and bbox)
|
||||
- imported: Uploaded file (requires file)
|
||||
"""
|
||||
|
||||
# Validate source_type
|
||||
if source_type not in ["ai_generated", "manual_selection", "imported"]:
|
||||
raise HTTPException(status_code=400, detail="Invalid source_type")
|
||||
|
||||
# Create patch record
|
||||
patch = Patch(
|
||||
name=name,
|
||||
description=description,
|
||||
source_type=source_type,
|
||||
source_project_id=source_project_id,
|
||||
source_edit_id=source_edit_id,
|
||||
tags=tags,
|
||||
category=category,
|
||||
file_path="", # Will be set after saving
|
||||
user_id=None # TODO: Add authentication
|
||||
)
|
||||
|
||||
db.add(patch)
|
||||
db.commit()
|
||||
db.refresh(patch)
|
||||
|
||||
# Save patch file based on source type
|
||||
patch_service = PatchLibraryService()
|
||||
|
||||
try:
|
||||
if source_type == "ai_generated":
|
||||
# Get edit directory and save AI-generated patch
|
||||
if not source_edit_id:
|
||||
raise HTTPException(status_code=400, detail="source_edit_id required for ai_generated")
|
||||
|
||||
edit = db.query(Edit).filter(Edit.id == source_edit_id).first()
|
||||
if not edit:
|
||||
raise HTTPException(status_code=404, detail="Edit not found")
|
||||
|
||||
from app.services.edit_service import EditService
|
||||
edit_service = EditService()
|
||||
edit_dir = edit_service.get_edit_dir(edit.project_id, edit.id)
|
||||
|
||||
file_path = patch_service.save_ai_generated_patch(patch.id, edit_dir)
|
||||
|
||||
# Get dimensions
|
||||
width, height = patch_service.get_patch_size(patch.id)
|
||||
patch.width = width
|
||||
patch.height = height
|
||||
|
||||
elif source_type == "manual_selection":
|
||||
# Save manually selected patch from current image
|
||||
if not source_project_id or not bbox:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="source_project_id and bbox required for manual_selection"
|
||||
)
|
||||
|
||||
project = db.query(Project).filter(Project.id == source_project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
bbox_dict = json.loads(bbox) if isinstance(bbox, str) else bbox
|
||||
file_path = patch_service.save_manual_patch(patch.id, source_project_id, bbox_dict)
|
||||
|
||||
patch.width = bbox_dict['width']
|
||||
patch.height = bbox_dict['height']
|
||||
|
||||
elif source_type == "imported":
|
||||
# Save uploaded file
|
||||
if not file:
|
||||
raise HTTPException(status_code=400, detail="file required for imported")
|
||||
|
||||
image_bytes = await file.read()
|
||||
file_path = patch_service.save_patch_from_bytes(patch.id, image_bytes)
|
||||
|
||||
# Get dimensions
|
||||
width, height = patch_service.get_patch_size(patch.id)
|
||||
patch.width = width
|
||||
patch.height = height
|
||||
|
||||
# Update patch with file path
|
||||
patch.file_path = file_path
|
||||
patch.thumbnail_path = str(patch_service.get_thumbnail_path(patch.id))
|
||||
db.commit()
|
||||
db.refresh(patch)
|
||||
|
||||
return patch
|
||||
|
||||
except Exception as e:
|
||||
# Cleanup on error
|
||||
patch_service.delete_patch(patch.id)
|
||||
db.delete(patch)
|
||||
db.commit()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/", response_model=List[PatchResponse])
|
||||
def list_patches(
|
||||
category: Optional[str] = None,
|
||||
tags: Optional[str] = None,
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""List patches in the library with optional filtering"""
|
||||
|
||||
query = db.query(Patch)
|
||||
|
||||
if category:
|
||||
query = query.filter(Patch.category == category)
|
||||
|
||||
if tags:
|
||||
# Simple tag search (could be improved with full-text search)
|
||||
query = query.filter(Patch.tags.like(f"%{tags}%"))
|
||||
|
||||
patches = query.offset(offset).limit(limit).all()
|
||||
return patches
|
||||
|
||||
|
||||
@router.get("/{patch_id}", response_model=PatchResponse)
|
||||
def get_patch(
|
||||
patch_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get patch details"""
|
||||
patch = db.query(Patch).filter(Patch.id == patch_id).first()
|
||||
if not patch:
|
||||
raise HTTPException(status_code=404, detail="Patch not found")
|
||||
return patch
|
||||
|
||||
|
||||
@router.get("/{patch_id}/image")
|
||||
def get_patch_image(
|
||||
patch_id: int,
|
||||
thumbnail: bool = False,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get patch image file"""
|
||||
patch = db.query(Patch).filter(Patch.id == patch_id).first()
|
||||
if not patch:
|
||||
raise HTTPException(status_code=404, detail="Patch not found")
|
||||
|
||||
patch_service = PatchLibraryService()
|
||||
|
||||
if thumbnail:
|
||||
file_path = patch_service.get_thumbnail_path(patch_id)
|
||||
else:
|
||||
file_path = patch_service.get_patch_path(patch_id)
|
||||
|
||||
if not file_path.exists():
|
||||
raise HTTPException(status_code=404, detail="Patch image not found")
|
||||
|
||||
return FileResponse(file_path, media_type="image/png")
|
||||
|
||||
|
||||
@router.post("/apply", response_model=StatusResponse)
|
||||
async def apply_patch(
|
||||
project_id: int = Form(...),
|
||||
patch_id: int = Form(...),
|
||||
bbox: str = Form(...),
|
||||
feather_px: int = Form(5),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Apply a saved patch to a project image
|
||||
|
||||
This creates a new edit in the project history.
|
||||
"""
|
||||
# Verify project exists
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Verify patch exists
|
||||
patch = db.query(Patch).filter(Patch.id == patch_id).first()
|
||||
if not patch:
|
||||
raise HTTPException(status_code=404, detail="Patch not found")
|
||||
|
||||
# Parse bbox
|
||||
bbox_dict = json.loads(bbox) if isinstance(bbox, str) else bbox
|
||||
|
||||
# Load current image
|
||||
from app.services.edit_service import EditService
|
||||
from PIL import Image
|
||||
|
||||
edit_service = EditService()
|
||||
current_image_path = edit_service.get_current_image_path(project_id)
|
||||
current_image = Image.open(current_image_path).convert('RGBA')
|
||||
|
||||
# Apply patch
|
||||
patch_service = PatchLibraryService()
|
||||
result_image = patch_service.apply_patch_to_image(
|
||||
patch_id,
|
||||
current_image,
|
||||
bbox_dict,
|
||||
feather_px
|
||||
)
|
||||
|
||||
# Save result as current image
|
||||
result_image.save(current_image_path)
|
||||
|
||||
# Create edit record
|
||||
edit = Edit(
|
||||
project_id=project_id,
|
||||
mode="patch_library",
|
||||
prompt=f"Applied saved patch: {patch.name}",
|
||||
selection_type="rectangle",
|
||||
bbox_json=json.dumps(bbox_dict),
|
||||
feather_px=feather_px,
|
||||
ai_provider="patch_library",
|
||||
status="completed"
|
||||
)
|
||||
db.add(edit)
|
||||
db.commit()
|
||||
|
||||
return StatusResponse(
|
||||
status="success",
|
||||
message=f"Applied patch '{patch.name}' to project",
|
||||
data={"edit_id": edit.id}
|
||||
)
|
||||
|
||||
|
||||
@router.delete("/{patch_id}", response_model=StatusResponse)
|
||||
def delete_patch(
|
||||
patch_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Delete a patch from the library"""
|
||||
patch = db.query(Patch).filter(Patch.id == patch_id).first()
|
||||
if not patch:
|
||||
raise HTTPException(status_code=404, detail="Patch not found")
|
||||
|
||||
# Delete files
|
||||
patch_service = PatchLibraryService()
|
||||
patch_service.delete_patch(patch_id)
|
||||
|
||||
# Delete record
|
||||
db.delete(patch)
|
||||
db.commit()
|
||||
|
||||
return StatusResponse(
|
||||
status="success",
|
||||
message=f"Deleted patch '{patch.name}'"
|
||||
)
|
||||
|
||||
|
||||
@router.put("/{patch_id}", response_model=PatchResponse)
|
||||
def update_patch(
|
||||
patch_id: int,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
category: Optional[str] = None,
|
||||
tags: Optional[str] = None,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Update patch metadata"""
|
||||
patch = db.query(Patch).filter(Patch.id == patch_id).first()
|
||||
if not patch:
|
||||
raise HTTPException(status_code=404, detail="Patch not found")
|
||||
|
||||
if name:
|
||||
patch.name = name
|
||||
if description is not None:
|
||||
patch.description = description
|
||||
if category:
|
||||
patch.category = category
|
||||
if tags is not None:
|
||||
patch.tags = tags
|
||||
|
||||
db.commit()
|
||||
db.refresh(patch)
|
||||
|
||||
return patch
|
||||
@@ -0,0 +1,487 @@
|
||||
"""
|
||||
Print / frame tools — frame fit and upscale.
|
||||
All endpoints under /api/print prefix.
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, Literal
|
||||
import base64
|
||||
import asyncio
|
||||
from io import BytesIO
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
|
||||
router = APIRouter(prefix="/api/print", tags=["print-tools"])
|
||||
|
||||
# ── Frame size catalogue (inches) ──────────────────────────────────────────
|
||||
FRAME_SIZES = {
|
||||
"4x6": (4, 6),
|
||||
"5x7": (5, 7),
|
||||
"8x10": (8, 10),
|
||||
"11x14": (11, 14),
|
||||
"16x20": (16, 20),
|
||||
"18x24": (18, 24),
|
||||
"20x24": (20, 24),
|
||||
"24x36": (24, 36),
|
||||
# Square
|
||||
"4x4": (4, 4),
|
||||
"8x8": (8, 8),
|
||||
"12x12": (12, 12),
|
||||
}
|
||||
|
||||
|
||||
def _encode(data: bytes) -> str:
|
||||
return base64.b64encode(data).decode()
|
||||
|
||||
|
||||
def _decode(b64: str) -> bytes:
|
||||
return base64.b64decode(b64)
|
||||
|
||||
|
||||
def _to_png(img: Image.Image) -> bytes:
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
# ── Request models ─────────────────────────────────────────────────────────
|
||||
|
||||
class FrameFitRequest(BaseModel):
|
||||
image: str # base64 PNG/JPEG
|
||||
frame: str # e.g. "8x10"
|
||||
orientation: Literal["auto", "portrait", "landscape"] = "auto"
|
||||
mode: Literal["crop", "extend", "smart"] = "smart"
|
||||
dpi: int = 300
|
||||
# For extend mode: prompt passed to outpaint
|
||||
prompt: Optional[str] = ""
|
||||
# Smart mode threshold: extend if gap fraction < this, else crop
|
||||
smart_threshold: float = 0.15
|
||||
|
||||
|
||||
class UpscaleRequest(BaseModel):
|
||||
image: str # base64
|
||||
scale: float = 2.0 # 1.5, 2, 3, 4
|
||||
# auto = pick best available; lanczos = always works; realesrgan_pytorch / realesrgan_ncnn = explicit
|
||||
method: str = "auto"
|
||||
|
||||
|
||||
class PrepareRequest(BaseModel):
|
||||
image: str # base64
|
||||
frame: str # e.g. "8x10"
|
||||
orientation: Literal["auto", "portrait", "landscape"] = "auto"
|
||||
target_dpi: int = 300
|
||||
upscale_method: str = "auto" # auto / realesrgan_pytorch / realesrgan_ncnn / lanczos
|
||||
mode: Literal["crop", "extend", "smart"] = "smart"
|
||||
prompt: Optional[str] = ""
|
||||
|
||||
|
||||
# ── Frame sizes endpoint ───────────────────────────────────────────────────
|
||||
|
||||
@router.get("/frame-sizes")
|
||||
def list_frame_sizes():
|
||||
"""Return the catalogue of supported frame sizes."""
|
||||
return {
|
||||
"sizes": list(FRAME_SIZES.keys()),
|
||||
"catalogue": {k: {"inches": v, "pixels_300dpi": (v[0]*300, v[1]*300)}
|
||||
for k, v in FRAME_SIZES.items()},
|
||||
}
|
||||
|
||||
|
||||
# ── Frame fit ──────────────────────────────────────────────────────────────
|
||||
|
||||
@router.post("/frame-fit")
|
||||
async def frame_fit(req: FrameFitRequest):
|
||||
"""
|
||||
Fit an image to a print frame size.
|
||||
|
||||
Modes:
|
||||
crop — center-crop to frame aspect ratio, then scale to print resolution.
|
||||
extend — scale to fill one dimension, outpaint the gap with AI.
|
||||
smart — extend if gap < smart_threshold of frame dimension, else crop.
|
||||
|
||||
Returns the fitted image plus a summary of what was done.
|
||||
"""
|
||||
if req.frame not in FRAME_SIZES:
|
||||
raise HTTPException(status_code=400,
|
||||
detail=f"Unknown frame '{req.frame}'. Valid: {list(FRAME_SIZES.keys())}")
|
||||
if not (72 <= req.dpi <= 600):
|
||||
raise HTTPException(status_code=400, detail="dpi must be 72–600")
|
||||
|
||||
try:
|
||||
image = Image.open(BytesIO(_decode(req.image))).convert("RGB")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Could not decode image: {e}")
|
||||
|
||||
fw, fh = FRAME_SIZES[req.frame] # frame inches (w, h in portrait)
|
||||
|
||||
# Resolve orientation
|
||||
img_w, img_h = image.size
|
||||
img_landscape = img_w >= img_h
|
||||
frame_landscape = fw >= fh
|
||||
|
||||
if req.orientation == "landscape":
|
||||
fw, fh = max(fw, fh), min(fw, fh)
|
||||
elif req.orientation == "portrait":
|
||||
fw, fh = min(fw, fh), max(fw, fh)
|
||||
else: # auto — match image orientation
|
||||
if img_landscape and not frame_landscape:
|
||||
fw, fh = fh, fw # rotate frame to landscape
|
||||
elif not img_landscape and frame_landscape:
|
||||
fw, fh = fh, fw # rotate frame to portrait
|
||||
|
||||
target_w = fw * req.dpi
|
||||
target_h = fh * req.dpi
|
||||
target_ratio = target_w / target_h
|
||||
img_ratio = img_w / img_h
|
||||
|
||||
# Determine actual mode
|
||||
mode = req.mode
|
||||
if mode == "smart":
|
||||
# Scale image to fill the frame — compute gap fraction
|
||||
if img_ratio > target_ratio:
|
||||
# Image wider → fits on height, gap on width
|
||||
scaled_h = target_h
|
||||
scaled_w = round(target_h * img_ratio)
|
||||
gap_frac = (scaled_w - target_w) / target_w # positive = overflow (crop)
|
||||
else:
|
||||
scaled_w = target_w
|
||||
scaled_h = round(target_w / img_ratio)
|
||||
gap_frac = (scaled_h - target_h) / target_h
|
||||
|
||||
# gap_frac > 0 means we'd need to crop; < 0 means we'd need to extend
|
||||
if gap_frac < 0:
|
||||
# Need to extend — use extend if gap is small enough
|
||||
mode = "extend" if abs(gap_frac) <= req.smart_threshold else "crop"
|
||||
else:
|
||||
mode = "crop"
|
||||
|
||||
if mode == "crop":
|
||||
result, summary = _crop_fit(image, target_w, target_h)
|
||||
else: # extend
|
||||
result, summary = await _extend_fit(image, target_w, target_h, req.prompt or "")
|
||||
|
||||
return {
|
||||
"result": _encode(_to_png(result)),
|
||||
"mode_used": mode,
|
||||
"frame": req.frame,
|
||||
"orientation": "landscape" if fw > fh else "portrait",
|
||||
"output_pixels": {"width": result.width, "height": result.height},
|
||||
"output_inches": {"width": fw, "height": fh},
|
||||
"dpi": req.dpi,
|
||||
"summary": summary,
|
||||
}
|
||||
|
||||
|
||||
def _crop_fit(image: Image.Image, target_w: int, target_h: int):
|
||||
"""Center-crop image to target aspect ratio, then Lanczos scale to target size."""
|
||||
img_w, img_h = image.size
|
||||
target_ratio = target_w / target_h
|
||||
img_ratio = img_w / img_h
|
||||
|
||||
if img_ratio > target_ratio:
|
||||
# Wider than target — crop sides
|
||||
new_w = round(img_h * target_ratio)
|
||||
x0 = (img_w - new_w) // 2
|
||||
cropped = image.crop((x0, 0, x0 + new_w, img_h))
|
||||
else:
|
||||
# Taller than target — crop top/bottom
|
||||
new_h = round(img_w / target_ratio)
|
||||
y0 = (img_h - new_h) // 2
|
||||
cropped = image.crop((0, y0, img_w, y0 + new_h))
|
||||
|
||||
result = cropped.resize((target_w, target_h), Image.Resampling.LANCZOS)
|
||||
summary = (
|
||||
f"Cropped from {img_w}×{img_h} to {cropped.width}×{cropped.height}, "
|
||||
f"scaled to {target_w}×{target_h}"
|
||||
)
|
||||
return result, summary
|
||||
|
||||
|
||||
async def _extend_fit(image: Image.Image, target_w: int, target_h: int, prompt: str):
|
||||
"""
|
||||
Scale image to fill one dimension exactly, then outpaint the gap with AI.
|
||||
Falls back to content-aware mirror fill if no remote provider configured.
|
||||
"""
|
||||
from app.services.remote_provider import get_remote_provider
|
||||
|
||||
img_w, img_h = image.size
|
||||
target_ratio = target_w / target_h
|
||||
img_ratio = img_w / img_h
|
||||
|
||||
if img_ratio > target_ratio:
|
||||
# Image wider — scale to target width, extend height
|
||||
scale = target_w / img_w
|
||||
scaled_w = target_w
|
||||
scaled_h = round(img_h * scale)
|
||||
gap_dir = "height"
|
||||
gap_top = (target_h - scaled_h) // 2
|
||||
gap_bottom = target_h - scaled_h - gap_top
|
||||
else:
|
||||
# Image taller — scale to target height, extend width
|
||||
scale = target_h / img_h
|
||||
scaled_h = target_h
|
||||
scaled_w = round(img_w * scale)
|
||||
gap_dir = "width"
|
||||
gap_left = (target_w - scaled_w) // 2
|
||||
gap_right = target_w - scaled_w - gap_left
|
||||
|
||||
scaled = image.resize((scaled_w, scaled_h), Image.Resampling.LANCZOS)
|
||||
|
||||
# Place scaled image on canvas
|
||||
canvas = Image.new("RGB", (target_w, target_h), (128, 128, 128))
|
||||
if gap_dir == "height":
|
||||
canvas.paste(scaled, (0, gap_top))
|
||||
# Build mask: top and bottom strips are white (to inpaint)
|
||||
mask = Image.new("L", (target_w, target_h), 0)
|
||||
if gap_top > 0:
|
||||
mask.paste(Image.new("L", (target_w, gap_top), 255), (0, 0))
|
||||
if gap_bottom > 0:
|
||||
mask.paste(Image.new("L", (target_w, gap_bottom), 255), (0, target_h - gap_bottom))
|
||||
else:
|
||||
canvas.paste(scaled, (gap_left, 0))
|
||||
mask = Image.new("L", (target_w, target_h), 0)
|
||||
if gap_left > 0:
|
||||
mask.paste(Image.new("L", (gap_left, target_h), 255), (0, 0))
|
||||
if gap_right > 0:
|
||||
mask.paste(Image.new("L", (gap_right, target_h), 255), (target_w - gap_right, 0))
|
||||
|
||||
# Try AI inpaint
|
||||
provider = get_remote_provider("inpaint")
|
||||
if provider:
|
||||
try:
|
||||
canvas_bytes = _to_png(canvas)
|
||||
mask_bytes = _to_png(mask)
|
||||
fill_prompt = prompt or "seamlessly continue the image, natural extension"
|
||||
result_bytes = await provider.inpaint(canvas_bytes, mask_bytes, fill_prompt, {})
|
||||
result = Image.open(BytesIO(result_bytes)).convert("RGB")
|
||||
summary = (
|
||||
f"Scaled {img_w}×{img_h} → {scaled_w}×{scaled_h}, "
|
||||
f"AI-extended {gap_dir} to {target_w}×{target_h}"
|
||||
)
|
||||
return result, summary
|
||||
except Exception as e:
|
||||
print(f"AI extend failed, using mirror fill: {e}")
|
||||
|
||||
# Fallback: mirror-fill the gap (looks decent for backgrounds/landscapes)
|
||||
result = _mirror_fill(canvas, mask, scaled, gap_dir,
|
||||
gap_top if gap_dir == "height" else gap_left,
|
||||
gap_bottom if gap_dir == "height" else gap_right,
|
||||
target_w, target_h)
|
||||
summary = (
|
||||
f"Scaled {img_w}×{img_h} → {scaled_w}×{scaled_h}, "
|
||||
f"mirror-filled {gap_dir} to {target_w}×{target_h} (no AI provider)"
|
||||
)
|
||||
return result, summary
|
||||
|
||||
|
||||
def _mirror_fill(canvas, mask, scaled, gap_dir, gap_a, gap_b, target_w, target_h):
|
||||
"""Fill gaps by reflecting the nearest edge strip."""
|
||||
result = canvas.copy()
|
||||
if gap_dir == "height":
|
||||
if gap_a > 0:
|
||||
strip = scaled.crop((0, 0, scaled.width, min(gap_a * 2, scaled.height)))
|
||||
strip = strip.transpose(Image.Transpose.FLIP_TOP_BOTTOM)
|
||||
strip = strip.resize((target_w, gap_a), Image.Resampling.LANCZOS)
|
||||
result.paste(strip, (0, 0))
|
||||
if gap_b > 0:
|
||||
strip = scaled.crop((0, max(0, scaled.height - gap_b * 2), scaled.width, scaled.height))
|
||||
strip = strip.transpose(Image.Transpose.FLIP_TOP_BOTTOM)
|
||||
strip = strip.resize((target_w, gap_b), Image.Resampling.LANCZOS)
|
||||
result.paste(strip, (0, target_h - gap_b))
|
||||
else:
|
||||
if gap_a > 0:
|
||||
strip = scaled.crop((0, 0, min(gap_a * 2, scaled.width), scaled.height))
|
||||
strip = strip.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
|
||||
strip = strip.resize((gap_a, target_h), Image.Resampling.LANCZOS)
|
||||
result.paste(strip, (0, 0))
|
||||
if gap_b > 0:
|
||||
strip = scaled.crop((max(0, scaled.width - gap_b * 2), 0, scaled.width, scaled.height))
|
||||
strip = strip.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
|
||||
strip = strip.resize((gap_b, target_h), Image.Resampling.LANCZOS)
|
||||
result.paste(strip, (target_w - gap_b, 0))
|
||||
return result
|
||||
|
||||
|
||||
# ── Upscale ────────────────────────────────────────────────────────────────
|
||||
|
||||
@router.post("/upscale/refresh-caps")
|
||||
def upscale_refresh_caps():
|
||||
"""Bust the capability cache (call after installing Real-ESRGAN without restarting)."""
|
||||
from app.services.upscale import invalidate_caps_cache, probe_upscale_capabilities
|
||||
invalidate_caps_cache()
|
||||
return probe_upscale_capabilities()
|
||||
|
||||
|
||||
@router.get("/upscale/available")
|
||||
async def upscale_available():
|
||||
"""
|
||||
Return capability probe: which upscale methods are available,
|
||||
which device will be used, and which method is recommended.
|
||||
If no AI upscaler is found, triggers background NCNN auto-install.
|
||||
Frontend uses this to populate the method selector.
|
||||
"""
|
||||
from app.services.upscale import probe_upscale_capabilities, ensure_ncnn_installed, get_install_status
|
||||
caps = probe_upscale_capabilities()
|
||||
# Auto-install NCNN if no AI upscaler is available yet
|
||||
if not caps["realesrgan_pytorch"] and not caps["realesrgan_ncnn"]:
|
||||
asyncio.create_task(ensure_ncnn_installed())
|
||||
caps["ncnn_install_status"] = get_install_status()
|
||||
return caps
|
||||
|
||||
|
||||
@router.get("/upscale/install-status")
|
||||
def upscale_install_status():
|
||||
"""Poll for Real-ESRGAN NCNN auto-install progress."""
|
||||
from app.services.upscale import get_install_status, probe_upscale_capabilities, _find_ncnn_binary
|
||||
status = get_install_status()
|
||||
# If install just finished, refresh caps
|
||||
if status["state"] == "done":
|
||||
from app.services.upscale import invalidate_caps_cache
|
||||
invalidate_caps_cache()
|
||||
caps = probe_upscale_capabilities()
|
||||
status["ncnn_available"] = caps["realesrgan_ncnn"]
|
||||
else:
|
||||
status["ncnn_available"] = False
|
||||
return status
|
||||
|
||||
|
||||
@router.post("/prepare")
|
||||
async def prepare_for_print(req: PrepareRequest):
|
||||
"""
|
||||
One-shot Prepare for Print: AI upscale to reach target DPI, then fit to frame.
|
||||
|
||||
Steps:
|
||||
1. Resolve target pixel dimensions (frame × target_dpi, orientation-adjusted)
|
||||
2. Calculate needed upscale factor so the image meets the target resolution
|
||||
3. Run Real-ESRGAN if scale > 1.05 (else skip — already large enough)
|
||||
4. Run frame-fit (crop / extend / smart) to exact target dimensions
|
||||
5. Return the print-ready image and a quality report
|
||||
"""
|
||||
if req.frame not in FRAME_SIZES:
|
||||
raise HTTPException(status_code=400,
|
||||
detail=f"Unknown frame '{req.frame}'. Valid: {list(FRAME_SIZES.keys())}")
|
||||
if not (72 <= req.target_dpi <= 600):
|
||||
raise HTTPException(status_code=400, detail="target_dpi must be 72–600")
|
||||
|
||||
try:
|
||||
image = Image.open(BytesIO(_decode(req.image))).convert("RGB")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Could not decode image: {e}")
|
||||
|
||||
fw, fh = FRAME_SIZES[req.frame]
|
||||
img_w, img_h = image.size
|
||||
|
||||
# Resolve orientation (same logic as frame_fit)
|
||||
img_landscape = img_w >= img_h
|
||||
frame_landscape = fw >= fh
|
||||
if req.orientation == "landscape":
|
||||
fw, fh = max(fw, fh), min(fw, fh)
|
||||
elif req.orientation == "portrait":
|
||||
fw, fh = min(fw, fh), max(fw, fh)
|
||||
else:
|
||||
if img_landscape and not frame_landscape:
|
||||
fw, fh = fh, fw
|
||||
elif not img_landscape and frame_landscape:
|
||||
fw, fh = fh, fw
|
||||
|
||||
target_w = fw * req.target_dpi
|
||||
target_h = fh * req.target_dpi
|
||||
|
||||
# Scale factor needed so the shorter dimension fills the frame
|
||||
scale_w = target_w / img_w
|
||||
scale_h = target_h / img_h
|
||||
needed_scale = min(scale_w, scale_h) # fill-to-fit (extend) baseline
|
||||
# For crop mode we need max; use the larger to be safe and let frame-fit crop
|
||||
needed_scale_crop = max(scale_w, scale_h)
|
||||
|
||||
# Use the smaller (extend) scale as the upscale target; frame-fit handles the rest
|
||||
upscale_factor = max(1.0, needed_scale)
|
||||
upscale_applied = False
|
||||
method_used = "none"
|
||||
|
||||
upscaled = image
|
||||
if upscale_factor > 1.05:
|
||||
# Cap per-pass at 4× (Real-ESRGAN works best at 2–4×)
|
||||
remaining = upscale_factor
|
||||
while remaining > 1.05:
|
||||
pass_scale = min(remaining, 4.0)
|
||||
# Round to one decimal to keep scale in 1.1–8.0 range accepted by upscale service
|
||||
pass_scale = round(pass_scale, 1)
|
||||
if pass_scale < 1.1:
|
||||
break
|
||||
from app.services.upscale import upscale_image
|
||||
result_bytes, method_used = await upscale_image(upscaled, pass_scale, req.upscale_method)
|
||||
upscaled = Image.open(BytesIO(result_bytes)).convert("RGB")
|
||||
remaining /= pass_scale
|
||||
upscale_applied = True
|
||||
|
||||
# Encode upscaled image and run frame-fit
|
||||
upscaled_b64 = _encode(_to_png(upscaled))
|
||||
|
||||
fit_req = FrameFitRequest(
|
||||
image=upscaled_b64,
|
||||
frame=req.frame,
|
||||
orientation=req.orientation,
|
||||
mode=req.mode,
|
||||
dpi=req.target_dpi,
|
||||
prompt=req.prompt or "",
|
||||
)
|
||||
# Re-use the existing frame_fit logic inline
|
||||
fit_response = await frame_fit(fit_req)
|
||||
|
||||
return {
|
||||
"result": fit_response["result"],
|
||||
"frame": req.frame,
|
||||
"orientation": fit_response["orientation"],
|
||||
"output_pixels": fit_response["output_pixels"],
|
||||
"output_inches": fit_response["output_inches"],
|
||||
"dpi": req.target_dpi,
|
||||
"mode_used": fit_response["mode_used"],
|
||||
"upscale_applied": upscale_applied,
|
||||
"upscale_factor": round(upscale_factor, 2),
|
||||
"upscale_method": method_used,
|
||||
"summary": fit_response["summary"],
|
||||
}
|
||||
|
||||
|
||||
@router.post("/upscale")
|
||||
async def upscale(req: UpscaleRequest):
|
||||
"""
|
||||
Upscale image. method values:
|
||||
auto — pick best available (recommended)
|
||||
realesrgan_pytorch — Real-ESRGAN via PyTorch (CUDA/MPS/CPU)
|
||||
realesrgan_ncnn — Real-ESRGAN NCNN Vulkan binary
|
||||
lanczos — always available, instant
|
||||
Any AI method falls back to the next best if unavailable.
|
||||
"""
|
||||
if not (1.1 <= req.scale <= 8.0):
|
||||
raise HTTPException(status_code=400, detail="scale must be 1.1–8.0")
|
||||
|
||||
valid_methods = {"auto", "realesrgan_pytorch", "realesrgan_ncnn", "lanczos"}
|
||||
if req.method not in valid_methods:
|
||||
raise HTTPException(status_code=400,
|
||||
detail=f"method must be one of {sorted(valid_methods)}")
|
||||
|
||||
try:
|
||||
image = Image.open(BytesIO(_decode(req.image))).convert("RGB")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Could not decode image: {e}")
|
||||
|
||||
orig_w, orig_h = image.size
|
||||
|
||||
try:
|
||||
from app.services.upscale import upscale_image
|
||||
result_bytes, method_used = await upscale_image(image, req.scale, req.method)
|
||||
result = Image.open(BytesIO(result_bytes))
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
return {
|
||||
"result": _encode(result_bytes),
|
||||
"method": method_used,
|
||||
"original": {"width": orig_w, "height": orig_h},
|
||||
"output": {"width": result.width, "height": result.height},
|
||||
"scale": req.scale,
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, status
|
||||
from sqlalchemy.orm import Session
|
||||
from typing import List
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from PIL import Image
|
||||
|
||||
from app.database import get_db
|
||||
from app.models.project import Project
|
||||
from app.models.edit import Edit
|
||||
from app.schemas import ProjectCreate, ProjectResponse, EditResponse, UploadResponse
|
||||
from app.services.edit_service import EditService
|
||||
from app.config import settings
|
||||
|
||||
router = APIRouter(prefix="/projects", tags=["projects"])
|
||||
|
||||
|
||||
@router.post("/", response_model=ProjectResponse)
|
||||
def create_project(
|
||||
project: ProjectCreate,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Create a new project"""
|
||||
# For MVP, we'll use a default user_id of 1
|
||||
# In production, this would come from authentication
|
||||
user_id = 1
|
||||
|
||||
db_project = Project(
|
||||
user_id=user_id,
|
||||
name=project.name
|
||||
)
|
||||
db.add(db_project)
|
||||
db.commit()
|
||||
db.refresh(db_project)
|
||||
|
||||
# Create project directory
|
||||
edit_service = EditService()
|
||||
edit_service.ensure_project_dir(db_project.id)
|
||||
|
||||
return db_project
|
||||
|
||||
|
||||
@router.get("/", response_model=List[ProjectResponse])
|
||||
def list_projects(
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""List all projects"""
|
||||
projects = db.query(Project).offset(skip).limit(limit).all()
|
||||
return projects
|
||||
|
||||
|
||||
@router.get("/{project_id}", response_model=ProjectResponse)
|
||||
def get_project(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get a specific project"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
return project
|
||||
|
||||
|
||||
@router.delete("/{project_id}")
|
||||
def delete_project(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Delete a project"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Delete project directory
|
||||
edit_service = EditService()
|
||||
project_dir = edit_service.get_project_dir(project_id)
|
||||
if project_dir.exists():
|
||||
shutil.rmtree(project_dir)
|
||||
|
||||
db.delete(project)
|
||||
db.commit()
|
||||
|
||||
return {"status": "success", "message": f"Project {project_id} deleted"}
|
||||
|
||||
|
||||
@router.post("/{project_id}/upload", response_model=UploadResponse)
|
||||
async def upload_image(
|
||||
project_id: int,
|
||||
file: UploadFile = File(...),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Upload an image to a project"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Validate file type
|
||||
if not file.content_type.startswith('image/'):
|
||||
raise HTTPException(status_code=400, detail="File must be an image")
|
||||
|
||||
# Create project directory
|
||||
edit_service = EditService()
|
||||
edit_service.ensure_project_dir(project_id)
|
||||
|
||||
# Save original and current images
|
||||
original_path = edit_service.get_original_image_path(project_id)
|
||||
current_path = edit_service.get_current_image_path(project_id)
|
||||
|
||||
# Read and validate image
|
||||
contents = await file.read()
|
||||
try:
|
||||
image = Image.open(BytesIO(contents))
|
||||
image = image.convert('RGBA')
|
||||
|
||||
# Save images
|
||||
image.save(original_path, 'PNG')
|
||||
image.save(current_path, 'PNG')
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Invalid image file: {str(e)}")
|
||||
|
||||
return UploadResponse(
|
||||
project_id=project_id,
|
||||
original_url=f"/projects/{project_id}/original",
|
||||
current_url=f"/projects/{project_id}/current"
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{project_id}/edits", response_model=List[EditResponse])
|
||||
def list_edits(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""List all edits for a project"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
edits = db.query(Edit).filter(Edit.project_id == project_id).order_by(Edit.created_at.desc()).all()
|
||||
return edits
|
||||
|
||||
|
||||
from io import BytesIO
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,138 @@
|
||||
from pydantic import BaseModel, EmailStr
|
||||
from typing import Optional, List, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
# User schemas
|
||||
class UserCreate(BaseModel):
|
||||
email: EmailStr
|
||||
password: str
|
||||
|
||||
|
||||
class UserResponse(BaseModel):
|
||||
id: int
|
||||
email: str
|
||||
created_at: datetime
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
# Project schemas
|
||||
class ProjectCreate(BaseModel):
|
||||
name: str
|
||||
|
||||
|
||||
class ProjectResponse(BaseModel):
|
||||
id: int
|
||||
user_id: int
|
||||
name: str
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
# Edit schemas
|
||||
class EditRequest(BaseModel):
|
||||
prompt: str
|
||||
mode: str = "A" # "A" or "B"
|
||||
selection_type: str # "rectangle", "ellipse", "lasso"
|
||||
bbox: Dict[str, int] # {x, y, width, height}
|
||||
feather_px: int = 0
|
||||
selection_data: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class EditResponse(BaseModel):
|
||||
id: int
|
||||
project_id: int
|
||||
created_at: datetime
|
||||
mode: str
|
||||
prompt: str
|
||||
selection_type: str
|
||||
bbox_json: str
|
||||
feather_px: int
|
||||
ai_provider: str
|
||||
status: str
|
||||
error_message: Optional[str] = None
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
# Image upload
|
||||
class UploadResponse(BaseModel):
|
||||
project_id: int
|
||||
original_url: str
|
||||
current_url: str
|
||||
|
||||
|
||||
# Patch Library schemas
|
||||
class PatchCreate(BaseModel):
|
||||
name: str
|
||||
description: Optional[str] = None
|
||||
source_type: str # "ai_generated", "manual_selection", "imported"
|
||||
source_project_id: Optional[int] = None
|
||||
source_edit_id: Optional[int] = None
|
||||
category: Optional[str] = None
|
||||
tags: Optional[str] = None
|
||||
bbox: Optional[Dict[str, int]] = None
|
||||
|
||||
|
||||
class PatchResponse(BaseModel):
|
||||
id: int
|
||||
name: str
|
||||
description: Optional[str]
|
||||
created_at: datetime
|
||||
source_type: str
|
||||
source_project_id: Optional[int]
|
||||
source_edit_id: Optional[int]
|
||||
width: int
|
||||
height: int
|
||||
tags: Optional[str]
|
||||
category: Optional[str]
|
||||
is_public: bool
|
||||
file_path: str
|
||||
thumbnail_path: Optional[str]
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
class PatchApply(BaseModel):
|
||||
project_id: int
|
||||
patch_id: int
|
||||
bbox: Dict[str, int]
|
||||
feather_px: int = 5
|
||||
|
||||
|
||||
# Text-to-Image schemas
|
||||
class TextToImageRequest(BaseModel):
|
||||
prompt: str
|
||||
width: int = 1024
|
||||
height: int = 1024
|
||||
negative_prompt: Optional[str] = None
|
||||
ai_provider: Optional[str] = None
|
||||
ai_model: Optional[str] = None
|
||||
create_project: bool = True
|
||||
project_name: Optional[str] = None
|
||||
|
||||
|
||||
class TextToImageResponse(BaseModel):
|
||||
status: str
|
||||
prompt: str
|
||||
width: int
|
||||
height: int
|
||||
project_id: Optional[int] = None
|
||||
image_url: Optional[str] = None
|
||||
layer_position: Optional[Dict[str, int]] = None
|
||||
ai_provider: str
|
||||
ai_model: Optional[str] = None
|
||||
|
||||
|
||||
# Generic responses
|
||||
class StatusResponse(BaseModel):
|
||||
status: str
|
||||
message: Optional[str] = None
|
||||
data: Optional[Any] = None
|
||||
@@ -0,0 +1,539 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Optional, Dict
|
||||
import httpx
|
||||
import base64
|
||||
import asyncio
|
||||
from io import BytesIO
|
||||
from app.config import settings
|
||||
|
||||
|
||||
class AIProvider(ABC):
|
||||
"""Abstract base class for AI providers"""
|
||||
|
||||
@abstractmethod
|
||||
async def edit_image(
|
||||
self,
|
||||
patch_image_bytes: bytes,
|
||||
mask_image_bytes: bytes,
|
||||
prompt: str,
|
||||
mode: str,
|
||||
full_image_bytes: Optional[bytes] = None,
|
||||
model: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""
|
||||
Edit an image patch using AI
|
||||
|
||||
Args:
|
||||
patch_image_bytes: The cropped patch to edit
|
||||
mask_image_bytes: Binary mask (same size as patch)
|
||||
prompt: Text description of desired changes
|
||||
mode: "A" (patch only) or "B" (patch + full image reference)
|
||||
full_image_bytes: Full image for context (mode B only)
|
||||
model: Optional specific model to use
|
||||
|
||||
Returns:
|
||||
Regenerated patch as bytes
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def text_to_image(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int = 1024,
|
||||
height: int = 1024,
|
||||
model: Optional[str] = None,
|
||||
negative_prompt: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""
|
||||
Generate an image from text prompt
|
||||
|
||||
Args:
|
||||
prompt: Text description of desired image
|
||||
width: Image width in pixels
|
||||
height: Image height in pixels
|
||||
model: Optional specific model to use
|
||||
negative_prompt: What to avoid in the generation
|
||||
|
||||
Returns:
|
||||
Generated image as bytes
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class OpenAIProvider(AIProvider):
|
||||
"""OpenAI DALL-E 2 based image editing (NOTE: Lower quality than DALL-E 3)"""
|
||||
|
||||
def __init__(self, api_key: str):
|
||||
self.api_key = api_key
|
||||
self.base_url = "https://api.openai.com/v1"
|
||||
|
||||
async def edit_image(
|
||||
self,
|
||||
patch_image_bytes: bytes,
|
||||
mask_image_bytes: bytes,
|
||||
prompt: str,
|
||||
mode: str,
|
||||
full_image_bytes: Optional[bytes] = None,
|
||||
model: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Edit image using OpenAI DALL-E 2 (NOTE: Uses older model, lower quality)"""
|
||||
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
files = {
|
||||
'image': ('image.png', patch_image_bytes, 'image/png'),
|
||||
'mask': ('mask.png', mask_image_bytes, 'image/png'),
|
||||
}
|
||||
|
||||
data = {
|
||||
'prompt': prompt,
|
||||
'n': 1,
|
||||
'size': '1024x1024' # Will be adjusted based on input
|
||||
}
|
||||
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}'
|
||||
}
|
||||
|
||||
response = await client.post(
|
||||
f"{self.base_url}/images/edits",
|
||||
files=files,
|
||||
data=data,
|
||||
headers=headers
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
# Download the generated image
|
||||
image_url = result['data'][0]['url']
|
||||
image_response = await client.get(image_url)
|
||||
image_response.raise_for_status()
|
||||
|
||||
return image_response.content
|
||||
|
||||
async def text_to_image(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int = 1024,
|
||||
height: int = 1024,
|
||||
model: Optional[str] = None,
|
||||
negative_prompt: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Generate image using OpenAI DALL-E"""
|
||||
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
data = {
|
||||
'prompt': prompt,
|
||||
'n': 1,
|
||||
'size': f'{width}x{height}' if width == height else '1024x1024'
|
||||
}
|
||||
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}'
|
||||
}
|
||||
|
||||
response = await client.post(
|
||||
f"{self.base_url}/images/generations",
|
||||
json=data,
|
||||
headers=headers
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
# Download the generated image
|
||||
image_url = result['data'][0]['url']
|
||||
image_response = await client.get(image_url)
|
||||
image_response.raise_for_status()
|
||||
|
||||
return image_response.content
|
||||
|
||||
|
||||
class StabilityAIProvider(AIProvider):
|
||||
"""Stability AI based image editing (SDXL Inpainting)"""
|
||||
|
||||
# Available Stability AI engines
|
||||
MODELS = {
|
||||
'sdxl': 'stable-diffusion-xl-1024-v1-0',
|
||||
'sd15': 'stable-diffusion-v1-5',
|
||||
'sd21': 'stable-diffusion-512-v2-1',
|
||||
}
|
||||
|
||||
def __init__(self, api_key: str, default_model: str = 'sdxl'):
|
||||
self.api_key = api_key
|
||||
self.base_url = "https://api.stability.ai/v1"
|
||||
self.default_model = default_model
|
||||
|
||||
async def edit_image(
|
||||
self,
|
||||
patch_image_bytes: bytes,
|
||||
mask_image_bytes: bytes,
|
||||
prompt: str,
|
||||
mode: str,
|
||||
full_image_bytes: Optional[bytes] = None,
|
||||
model: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Edit image using Stability AI SDXL Inpainting"""
|
||||
|
||||
# Select model
|
||||
model_key = model or self.default_model
|
||||
engine_id = self.MODELS.get(model_key, self.MODELS['sdxl'])
|
||||
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
files = {
|
||||
'init_image': ('image.png', patch_image_bytes, 'image/png'),
|
||||
'mask_image': ('mask.png', mask_image_bytes, 'image/png'),
|
||||
}
|
||||
|
||||
# Optimized parameters for better quality
|
||||
data = {
|
||||
'text_prompts[0][text]': prompt,
|
||||
'text_prompts[0][weight]': '1.0',
|
||||
'cfg_scale': '8', # Increased for better prompt adherence
|
||||
'samples': '1',
|
||||
'steps': '40', # Increased for better quality
|
||||
'mask_source': 'MASK_IMAGE_WHITE', # White areas are inpainted
|
||||
}
|
||||
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Accept': 'application/json'
|
||||
}
|
||||
|
||||
response = await client.post(
|
||||
f"{self.base_url}/generation/{engine_id}/image-to-image/masking",
|
||||
files=files,
|
||||
data=data,
|
||||
headers=headers
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
# Decode base64 image
|
||||
image_data = result['artifacts'][0]['base64']
|
||||
return base64.b64decode(image_data)
|
||||
|
||||
async def text_to_image(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int = 1024,
|
||||
height: int = 1024,
|
||||
model: Optional[str] = None,
|
||||
negative_prompt: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Generate image using Stability AI SDXL"""
|
||||
|
||||
# Select model
|
||||
model_key = model or self.default_model
|
||||
engine_id = self.MODELS.get(model_key, self.MODELS['sdxl'])
|
||||
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
# Build prompts array
|
||||
data = {
|
||||
'text_prompts[0][text]': prompt,
|
||||
'text_prompts[0][weight]': '1.0',
|
||||
'cfg_scale': '7',
|
||||
'samples': '1',
|
||||
'steps': '50',
|
||||
'height': str(height),
|
||||
'width': str(width),
|
||||
}
|
||||
|
||||
# Add negative prompt if provided
|
||||
if negative_prompt:
|
||||
data['text_prompts[1][text]'] = negative_prompt
|
||||
data['text_prompts[1][weight]'] = '-1.0'
|
||||
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Accept': 'application/json'
|
||||
}
|
||||
|
||||
response = await client.post(
|
||||
f"{self.base_url}/generation/{engine_id}/text-to-image",
|
||||
data=data,
|
||||
headers=headers
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
# Decode base64 image
|
||||
image_data = result['artifacts'][0]['base64']
|
||||
return base64.b64decode(image_data)
|
||||
|
||||
|
||||
class ReplicateProvider(AIProvider):
|
||||
"""Replicate API with multiple model support"""
|
||||
|
||||
# Available Replicate models for inpainting
|
||||
MODELS = {
|
||||
# SDXL Inpainting - Best general purpose
|
||||
'sdxl-inpaint': {
|
||||
'version': 'stability-ai/sdxl:39ed52f2a78e934b3ba6e2a89f5b1c712de7dfea535525255b1aa35c5565e08b',
|
||||
'use_case': 'General purpose, high quality',
|
||||
'cost': '~$0.025/image',
|
||||
'best_for': ['general', 'landscapes', 'objects', 'textures']
|
||||
},
|
||||
# LaMa - Best for object removal
|
||||
'lama': {
|
||||
'version': 'andreasjansson/lama:7f4a2e3c95ab83c1d66ea26a66c27f93b64a2e5a3c5f7f4f4f4f4f4f4f4f4f4f',
|
||||
'use_case': 'Object removal and cleanup',
|
||||
'cost': '~$0.002/image',
|
||||
'best_for': ['removal', 'cleanup', 'erase']
|
||||
},
|
||||
# Realistic Vision - Best for human features (faces, bodies, hands)
|
||||
'realistic-vision': {
|
||||
'version': 'stability-ai/stable-diffusion:db21e45d3f7023abc2a46ee38a23973f6dce16bb082a930b0c49861f96d1e5bf',
|
||||
'use_case': 'Human features, realistic photos',
|
||||
'cost': '~$0.020/image',
|
||||
'best_for': ['face', 'body', 'hands', 'portrait', 'person', 'human']
|
||||
},
|
||||
}
|
||||
|
||||
def __init__(self, api_key: str, default_model: str = 'sdxl-inpaint'):
|
||||
self.api_key = api_key
|
||||
self.base_url = "https://api.replicate.com/v1"
|
||||
self.default_model = default_model
|
||||
|
||||
def _select_model_from_prompt(self, prompt: str) -> str:
|
||||
"""Auto-select best model based on prompt keywords"""
|
||||
prompt_lower = prompt.lower()
|
||||
|
||||
# Check for removal/cleanup keywords
|
||||
if any(word in prompt_lower for word in ['remove', 'erase', 'delete', 'cleanup']):
|
||||
return 'lama'
|
||||
|
||||
# Check for human feature keywords
|
||||
if any(word in prompt_lower for word in ['hand', 'face', 'body', 'person', 'portrait', 'skin']):
|
||||
return 'realistic-vision'
|
||||
|
||||
# Default to SDXL for general purpose
|
||||
return 'sdxl-inpaint'
|
||||
|
||||
async def edit_image(
|
||||
self,
|
||||
patch_image_bytes: bytes,
|
||||
mask_image_bytes: bytes,
|
||||
prompt: str,
|
||||
mode: str,
|
||||
full_image_bytes: Optional[bytes] = None,
|
||||
model: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Edit image using Replicate with auto model selection"""
|
||||
|
||||
# Auto-select model if not specified
|
||||
if not model:
|
||||
model = self._select_model_from_prompt(prompt)
|
||||
|
||||
model_config = self.MODELS.get(model, self.MODELS['sdxl-inpaint'])
|
||||
|
||||
# Convert bytes to base64 for Replicate API
|
||||
patch_b64 = base64.b64encode(patch_image_bytes).decode('utf-8')
|
||||
mask_b64 = base64.b64encode(mask_image_bytes).decode('utf-8')
|
||||
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
# Create prediction
|
||||
prediction_data = {
|
||||
"version": model_config['version'],
|
||||
"input": {
|
||||
"image": f"data:image/png;base64,{patch_b64}",
|
||||
"mask": f"data:image/png;base64,{mask_b64}",
|
||||
"prompt": prompt,
|
||||
"num_outputs": 1,
|
||||
"guidance_scale": 7.5,
|
||||
"num_inference_steps": 50,
|
||||
}
|
||||
}
|
||||
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
# Start prediction
|
||||
response = await client.post(
|
||||
f"{self.base_url}/predictions",
|
||||
json=prediction_data,
|
||||
headers=headers
|
||||
)
|
||||
response.raise_for_status()
|
||||
prediction = response.json()
|
||||
|
||||
# Poll for completion
|
||||
prediction_url = prediction['urls']['get']
|
||||
max_attempts = 60 # 2 minutes max
|
||||
attempt = 0
|
||||
|
||||
while attempt < max_attempts:
|
||||
await asyncio.sleep(2) # Wait 2 seconds between polls
|
||||
|
||||
status_response = await client.get(prediction_url, headers=headers)
|
||||
status_response.raise_for_status()
|
||||
status_data = status_response.json()
|
||||
|
||||
if status_data['status'] == 'succeeded':
|
||||
# Download result image
|
||||
output_url = status_data['output'][0]
|
||||
image_response = await client.get(output_url)
|
||||
image_response.raise_for_status()
|
||||
return image_response.content
|
||||
|
||||
elif status_data['status'] == 'failed':
|
||||
raise Exception(f"Replicate prediction failed: {status_data.get('error')}")
|
||||
|
||||
attempt += 1
|
||||
|
||||
raise Exception("Replicate prediction timed out")
|
||||
|
||||
async def text_to_image(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int = 1024,
|
||||
height: int = 1024,
|
||||
model: Optional[str] = None,
|
||||
negative_prompt: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Generate image using Replicate SDXL"""
|
||||
|
||||
# Use SDXL for text-to-image
|
||||
model_version = 'stability-ai/sdxl:39ed52f2a78e934b3ba6e2a89f5b1c712de7dfea535525255b1aa35c5565e08b'
|
||||
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
# Create prediction
|
||||
prediction_data = {
|
||||
"version": model_version,
|
||||
"input": {
|
||||
"prompt": prompt,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"num_outputs": 1,
|
||||
"guidance_scale": 7.5,
|
||||
"num_inference_steps": 50,
|
||||
}
|
||||
}
|
||||
|
||||
# Add negative prompt if provided
|
||||
if negative_prompt:
|
||||
prediction_data["input"]["negative_prompt"] = negative_prompt
|
||||
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
# Start prediction
|
||||
response = await client.post(
|
||||
f"{self.base_url}/predictions",
|
||||
json=prediction_data,
|
||||
headers=headers
|
||||
)
|
||||
response.raise_for_status()
|
||||
prediction = response.json()
|
||||
|
||||
# Poll for completion
|
||||
prediction_url = prediction['urls']['get']
|
||||
max_attempts = 60
|
||||
attempt = 0
|
||||
|
||||
while attempt < max_attempts:
|
||||
await asyncio.sleep(2)
|
||||
|
||||
status_response = await client.get(prediction_url, headers=headers)
|
||||
status_response.raise_for_status()
|
||||
status_data = status_response.json()
|
||||
|
||||
if status_data['status'] == 'succeeded':
|
||||
# Download result image
|
||||
output_url = status_data['output'][0]
|
||||
image_response = await client.get(output_url)
|
||||
image_response.raise_for_status()
|
||||
return image_response.content
|
||||
|
||||
elif status_data['status'] == 'failed':
|
||||
raise Exception(f"Replicate prediction failed: {status_data.get('error')}")
|
||||
|
||||
attempt += 1
|
||||
|
||||
raise Exception("Replicate text-to-image timed out")
|
||||
|
||||
|
||||
class MockAIProvider(AIProvider):
|
||||
"""Mock provider for testing (returns original patch)"""
|
||||
|
||||
async def edit_image(
|
||||
self,
|
||||
patch_image_bytes: bytes,
|
||||
mask_image_bytes: bytes,
|
||||
prompt: str,
|
||||
mode: str,
|
||||
full_image_bytes: Optional[bytes] = None,
|
||||
model: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Return the original patch (for testing)"""
|
||||
return patch_image_bytes
|
||||
|
||||
async def text_to_image(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int = 1024,
|
||||
height: int = 1024,
|
||||
model: Optional[str] = None,
|
||||
negative_prompt: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Generate a placeholder image (for testing)"""
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
# Create a simple placeholder image
|
||||
img = Image.new('RGB', (width, height), color='lightgray')
|
||||
draw = ImageDraw.Draw(img)
|
||||
|
||||
# Draw text
|
||||
text = f"Mock Image\n{width}x{height}\n{prompt[:50]}"
|
||||
draw.text((width//4, height//2), text, fill='black')
|
||||
|
||||
# Convert to bytes
|
||||
buffer = BytesIO()
|
||||
img.save(buffer, format='PNG')
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def get_ai_provider(provider_name: Optional[str] = None, model: Optional[str] = None) -> AIProvider:
|
||||
"""
|
||||
Factory function to get the configured AI provider
|
||||
|
||||
Args:
|
||||
provider_name: Override default provider from settings
|
||||
model: Specific model to use (provider-dependent)
|
||||
|
||||
Returns:
|
||||
AIProvider instance
|
||||
"""
|
||||
|
||||
provider = provider_name or settings.ai_provider
|
||||
provider = provider.lower()
|
||||
|
||||
if provider == "openai":
|
||||
if not settings.openai_api_key:
|
||||
raise ValueError("OpenAI API key not configured")
|
||||
return OpenAIProvider(settings.openai_api_key)
|
||||
|
||||
elif provider == "stability":
|
||||
if not settings.stability_api_key:
|
||||
raise ValueError("Stability AI API key not configured")
|
||||
default_model = model or getattr(settings, 'stability_model', 'sdxl')
|
||||
return StabilityAIProvider(settings.stability_api_key, default_model=default_model)
|
||||
|
||||
elif provider == "replicate":
|
||||
if not settings.replicate_api_key:
|
||||
raise ValueError("Replicate API key not configured")
|
||||
default_model = model or getattr(settings, 'replicate_model', 'sdxl-inpaint')
|
||||
return ReplicateProvider(settings.replicate_api_key, default_model=default_model)
|
||||
|
||||
elif provider == "mock":
|
||||
return MockAIProvider()
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown AI provider: {provider}")
|
||||
@@ -0,0 +1,220 @@
|
||||
import os
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional
|
||||
from datetime import datetime
|
||||
from PIL import Image
|
||||
|
||||
from app.models.edit import Edit
|
||||
from app.models.project import Project
|
||||
from app.services.ai_provider import get_ai_provider
|
||||
from app.utils.image_processing import (
|
||||
bytes_to_image,
|
||||
image_to_bytes,
|
||||
crop_patch,
|
||||
blend_patch,
|
||||
insert_patch,
|
||||
create_mask_from_selection,
|
||||
resize_for_ai,
|
||||
scale_bbox
|
||||
)
|
||||
from app.config import settings
|
||||
|
||||
|
||||
class EditService:
|
||||
"""Service for handling image edits"""
|
||||
|
||||
def __init__(self, data_dir: str = None):
|
||||
self.data_dir = data_dir or settings.data_dir
|
||||
self.ai_provider = get_ai_provider()
|
||||
|
||||
def get_project_dir(self, project_id: int) -> Path:
|
||||
"""Get project directory path"""
|
||||
return Path(self.data_dir) / "projects" / str(project_id)
|
||||
|
||||
def get_edit_dir(self, project_id: int, edit_id: int) -> Path:
|
||||
"""Get edit history directory path"""
|
||||
return self.get_project_dir(project_id) / "history" / str(edit_id)
|
||||
|
||||
def ensure_project_dir(self, project_id: int):
|
||||
"""Ensure project directory structure exists"""
|
||||
project_dir = self.get_project_dir(project_id)
|
||||
project_dir.mkdir(parents=True, exist_ok=True)
|
||||
(project_dir / "history").mkdir(exist_ok=True)
|
||||
|
||||
def get_current_image_path(self, project_id: int) -> Path:
|
||||
"""Get path to current image"""
|
||||
return self.get_project_dir(project_id) / "current.png"
|
||||
|
||||
def get_original_image_path(self, project_id: int) -> Path:
|
||||
"""Get path to original image"""
|
||||
return self.get_project_dir(project_id) / "original.png"
|
||||
|
||||
async def process_edit(
|
||||
self,
|
||||
project_id: int,
|
||||
edit_id: int,
|
||||
prompt: str,
|
||||
mode: str,
|
||||
selection_type: str,
|
||||
bbox: Dict[str, int],
|
||||
feather_px: int,
|
||||
selection_data: Optional[Dict] = None
|
||||
) -> str:
|
||||
"""
|
||||
Process an edit request
|
||||
|
||||
Args:
|
||||
project_id: Project ID
|
||||
edit_id: Edit ID
|
||||
prompt: AI prompt
|
||||
mode: "A" or "B"
|
||||
selection_type: "rectangle", "ellipse", or "lasso"
|
||||
bbox: Bounding box {x, y, width, height}
|
||||
feather_px: Feather radius in pixels
|
||||
selection_data: Additional selection data (for lasso)
|
||||
|
||||
Returns:
|
||||
Path to the result image
|
||||
"""
|
||||
# Create edit directory
|
||||
edit_dir = self.get_edit_dir(project_id, edit_id)
|
||||
edit_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Load current image
|
||||
current_image_path = self.get_current_image_path(project_id)
|
||||
full_image = Image.open(current_image_path).convert('RGBA')
|
||||
|
||||
# Crop patch from current image
|
||||
original_patch = crop_patch(full_image, bbox)
|
||||
|
||||
# Save original patch
|
||||
original_patch.save(edit_dir / "patch_in.png")
|
||||
|
||||
# Create mask based on selection type
|
||||
mask = create_mask_from_selection(
|
||||
bbox['width'],
|
||||
bbox['height'],
|
||||
selection_type,
|
||||
selection_data or {}
|
||||
)
|
||||
|
||||
# Save mask
|
||||
mask.save(edit_dir / "mask.png")
|
||||
|
||||
# Resize patch and mask for AI if needed
|
||||
patch_for_ai, scale = resize_for_ai(original_patch)
|
||||
mask_for_ai = mask.resize(patch_for_ai.size, Image.Resampling.LANCZOS)
|
||||
|
||||
# Prepare full image for mode B
|
||||
full_image_bytes = None
|
||||
if mode == "B":
|
||||
full_image_for_ai, _ = resize_for_ai(full_image)
|
||||
full_image_bytes = image_to_bytes(full_image_for_ai)
|
||||
|
||||
# Call AI provider
|
||||
regenerated_patch_bytes = await self.ai_provider.edit_image(
|
||||
patch_image_bytes=image_to_bytes(patch_for_ai),
|
||||
mask_image_bytes=image_to_bytes(mask_for_ai),
|
||||
prompt=prompt,
|
||||
mode=mode,
|
||||
full_image_bytes=full_image_bytes
|
||||
)
|
||||
|
||||
# Convert regenerated patch back to PIL Image
|
||||
regenerated_patch = bytes_to_image(regenerated_patch_bytes)
|
||||
|
||||
# Resize back to original patch size if scaled
|
||||
if scale != 1.0:
|
||||
regenerated_patch = regenerated_patch.resize(
|
||||
original_patch.size,
|
||||
Image.Resampling.LANCZOS
|
||||
)
|
||||
|
||||
# Save regenerated patch
|
||||
regenerated_patch.save(edit_dir / "patch_out.png")
|
||||
|
||||
# Blend regenerated patch with original using mask
|
||||
blended_patch = blend_patch(
|
||||
original_patch,
|
||||
regenerated_patch,
|
||||
mask,
|
||||
feather_px
|
||||
)
|
||||
|
||||
# Insert blended patch back into full image
|
||||
result_image = insert_patch(full_image, blended_patch, bbox)
|
||||
|
||||
# Save result
|
||||
result_path = edit_dir / "result.png"
|
||||
result_image.save(result_path)
|
||||
|
||||
# Update current image
|
||||
result_image.save(current_image_path)
|
||||
|
||||
# Save metadata
|
||||
metadata = {
|
||||
'edit_id': edit_id,
|
||||
'project_id': project_id,
|
||||
'prompt': prompt,
|
||||
'mode': mode,
|
||||
'selection_type': selection_type,
|
||||
'bbox': bbox,
|
||||
'feather_px': feather_px,
|
||||
'selection_data': selection_data,
|
||||
'timestamp': datetime.utcnow().isoformat(),
|
||||
'ai_provider': settings.ai_provider
|
||||
}
|
||||
|
||||
with open(edit_dir / "meta.json", 'w') as f:
|
||||
json.dump(metadata, f, indent=2)
|
||||
|
||||
return str(result_path)
|
||||
|
||||
def revert_to_edit(self, project_id: int, edit_id: int) -> str:
|
||||
"""
|
||||
Revert project to a specific edit
|
||||
|
||||
Args:
|
||||
project_id: Project ID
|
||||
edit_id: Edit ID to revert to
|
||||
|
||||
Returns:
|
||||
Path to the reverted image
|
||||
"""
|
||||
edit_dir = self.get_edit_dir(project_id, edit_id)
|
||||
result_path = edit_dir / "result.png"
|
||||
|
||||
if not result_path.exists():
|
||||
raise FileNotFoundError(f"Edit {edit_id} result not found")
|
||||
|
||||
# Copy result to current (preserve alpha channel)
|
||||
current_path = self.get_current_image_path(project_id)
|
||||
img = Image.open(result_path)
|
||||
# Preserve original mode to maintain transparency
|
||||
img.save(current_path, format='PNG')
|
||||
|
||||
return str(current_path)
|
||||
|
||||
def reset_to_original(self, project_id: int) -> str:
|
||||
"""
|
||||
Reset project to original image
|
||||
|
||||
Args:
|
||||
project_id: Project ID
|
||||
|
||||
Returns:
|
||||
Path to the original image
|
||||
"""
|
||||
original_path = self.get_original_image_path(project_id)
|
||||
current_path = self.get_current_image_path(project_id)
|
||||
|
||||
if not original_path.exists():
|
||||
raise FileNotFoundError(f"Original image for project {project_id} not found")
|
||||
|
||||
# Copy original to current (preserve alpha channel)
|
||||
img = Image.open(original_path)
|
||||
# Preserve original mode to maintain transparency
|
||||
img.save(current_path, format='PNG')
|
||||
|
||||
return str(current_path)
|
||||
@@ -0,0 +1,391 @@
|
||||
"""
|
||||
GPU capability detection and per-operation model selection.
|
||||
|
||||
Probes the actual GPU — VRAM (total + free), CUDA compute capability, and
|
||||
feature flags (fp16, bf16, fp8, int8, tensor cores) — then selects the
|
||||
highest-quality model that fits for each operation.
|
||||
|
||||
Model selection ladder (txt2img):
|
||||
eff_vram ≥ 20 GB → FLUX.1-schnell (no offload)
|
||||
eff_vram ≥ 10 GB → FLUX.1-schnell (model_cpu_offload, 2–3× slower but fits)
|
||||
eff_vram ≥ 7.5 GB → SDXL base
|
||||
eff_vram ≥ 5.5 GB → SDXL base + attention slicing
|
||||
eff_vram ≥ 4.0 GB → SDXL + model_cpu_offload (GTX 1060 6 GB, Quadro 6 GB)
|
||||
eff_vram ≥ 3.5 GB → Stable Diffusion 2.1
|
||||
eff_vram ≥ 2.5 GB → SD 2.1-base + attention slicing
|
||||
eff_vram ≥ 1.7 GB → Stable Diffusion 1.5
|
||||
otherwise → SD 1.5 + sequential CPU offload
|
||||
|
||||
Inpaint always uses SDXL/SD-family (no FLUX inpaint pipeline yet).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
|
||||
# ── Model specification ───────────────────────────────────────────────────────
|
||||
|
||||
@dataclass
|
||||
class ModelSpec:
|
||||
"""Everything needed to load and run one diffusion pipeline."""
|
||||
model_id: str
|
||||
family: str # sd15 | sd2x | sdxl | flux
|
||||
memory_opt: str # none | attention_slicing | model_cpu_offload | sequential_cpu_offload
|
||||
native_res: int # 512 | 768 | 1024
|
||||
vram_fp16_gb: float # approx VRAM needed in fp16, no memory opts
|
||||
|
||||
|
||||
# ── GPU capability record ─────────────────────────────────────────────────────
|
||||
|
||||
@dataclass
|
||||
class GpuCapabilities:
|
||||
# Hardware
|
||||
backend: str # cuda | mps | cpu
|
||||
device_name: str
|
||||
vram_total_gb: float
|
||||
vram_free_gb: float
|
||||
compute_capability: str # "8.6", "7.5", "6.1" …
|
||||
cc_major: int
|
||||
cc_minor: int
|
||||
|
||||
# Feature flags derived from compute capability
|
||||
fp16: bool # reliable fp16 (CC ≥ 6.0; CC 5.x works but slower)
|
||||
bf16: bool # native bf16 (CC ≥ 8.0)
|
||||
fp8: bool # native fp8 (CC ≥ 8.9, Ada / Hopper)
|
||||
int8: bool # efficient int8 (CC ≥ 7.0, needed for bitsandbytes)
|
||||
tensor_cores: bool # tensor cores (CC ≥ 7.0, Volta+)
|
||||
xformers: bool # xformers installed (reduces attention VRAM ~20-30%)
|
||||
|
||||
# Derived budget
|
||||
effective_vram_gb: float # free VRAM after overhead, halved if fp32-only
|
||||
|
||||
# Human-readable tier label
|
||||
tier: str # flux_full | flux_offload | sdxl | sdxl_low | sdxl_offload | sd2x | sd2x_low | sd15 | minimal
|
||||
|
||||
# Best model per operation
|
||||
recommended: dict[str, Optional[ModelSpec]]
|
||||
|
||||
# Metadata
|
||||
warnings: list[str]
|
||||
capabilities: list[str]
|
||||
|
||||
|
||||
# ── Detection ─────────────────────────────────────────────────────────────────
|
||||
|
||||
def detect_gpu() -> GpuCapabilities:
|
||||
"""Probe the GPU, return a fully populated GpuCapabilities."""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
props = torch.cuda.get_device_properties(0)
|
||||
free_bytes, total_bytes = torch.cuda.mem_get_info(0)
|
||||
vram_total = total_bytes / (1024 ** 3)
|
||||
vram_free = free_bytes / (1024 ** 3)
|
||||
cc = f"{props.major}.{props.minor}"
|
||||
major, minor = props.major, props.minor
|
||||
|
||||
fp16 = major >= 6 # Pascal and newer have good fp16
|
||||
bf16 = major >= 8 # Ampere A100 / RTX 3000+
|
||||
fp8 = major > 8 or (major == 8 and minor >= 9) # Ada / Hopper
|
||||
int8 = major >= 7 # Volta+
|
||||
tensor_cores = major >= 7
|
||||
|
||||
# Pre-Pascal (Maxwell CC 5.x): fp16 works but throughput is lower than fp32
|
||||
# on some Maxwell cards. Flag it so memory opt logic can account for it.
|
||||
xf = _xformers_available()
|
||||
|
||||
# Subtract driver/CUDA context overhead from free VRAM
|
||||
overhead_gb = 0.4
|
||||
eff = max(0.0, vram_free - overhead_gb)
|
||||
if not fp16:
|
||||
eff /= 2.0 # fp32 weights are 2× larger
|
||||
|
||||
tier = _tier_label(eff)
|
||||
warnings = _build_warnings(
|
||||
tier, vram_total, vram_free, cc, major, minor, fp16, bf16, fp8, xf
|
||||
)
|
||||
|
||||
return GpuCapabilities(
|
||||
backend="cuda",
|
||||
device_name=props.name,
|
||||
vram_total_gb=round(vram_total, 1),
|
||||
vram_free_gb=round(vram_free, 1),
|
||||
compute_capability=cc,
|
||||
cc_major=major,
|
||||
cc_minor=minor,
|
||||
fp16=fp16,
|
||||
bf16=bf16,
|
||||
fp8=fp8,
|
||||
int8=int8,
|
||||
tensor_cores=tensor_cores,
|
||||
xformers=xf,
|
||||
effective_vram_gb=round(eff, 1),
|
||||
tier=tier,
|
||||
recommended=_select_all_models(eff),
|
||||
warnings=warnings,
|
||||
capabilities=_caps(tier),
|
||||
)
|
||||
|
||||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
usable_gb = _apple_usable_gb()
|
||||
eff = max(0.0, usable_gb - 0.5)
|
||||
tier = _tier_label(eff)
|
||||
return GpuCapabilities(
|
||||
backend="mps",
|
||||
device_name="Apple Silicon",
|
||||
vram_total_gb=round(usable_gb, 1),
|
||||
vram_free_gb=round(usable_gb, 1),
|
||||
compute_capability="mps",
|
||||
cc_major=0,
|
||||
cc_minor=0,
|
||||
fp16=False, # MPS diffusion more stable in fp32
|
||||
bf16=False,
|
||||
fp8=False,
|
||||
int8=False,
|
||||
tensor_cores=False,
|
||||
xformers=False,
|
||||
effective_vram_gb=round(eff / 2, 1), # fp32 on MPS
|
||||
tier=tier,
|
||||
recommended=_select_all_models(eff / 2),
|
||||
warnings=["Apple MPS: using fp32 (fp16 less stable). Models load slower."],
|
||||
capabilities=_caps(tier),
|
||||
)
|
||||
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# CPU fallback
|
||||
return GpuCapabilities(
|
||||
backend="cpu",
|
||||
device_name="CPU (no GPU)",
|
||||
vram_total_gb=0.0,
|
||||
vram_free_gb=0.0,
|
||||
compute_capability="",
|
||||
cc_major=0, cc_minor=0,
|
||||
fp16=False, bf16=False, fp8=False, int8=False,
|
||||
tensor_cores=False, xformers=False,
|
||||
effective_vram_gb=0.0,
|
||||
tier="minimal",
|
||||
recommended=_select_all_models(0.0),
|
||||
warnings=[
|
||||
"No GPU found. Running on CPU — expect 5–30 minutes per image. "
|
||||
"Consider setting AI_PROVIDER to a remote/cloud provider instead."
|
||||
],
|
||||
capabilities=["txt2img", "inpaint", "img2img", "outpaint"],
|
||||
)
|
||||
|
||||
|
||||
# ── Model selection ───────────────────────────────────────────────────────────
|
||||
|
||||
def _select_all_models(eff_vram: float) -> dict[str, Optional[ModelSpec]]:
|
||||
return {
|
||||
"txt2img": _select_txt2img(eff_vram),
|
||||
"img2img": _select_img2img(eff_vram),
|
||||
"inpaint": _select_inpaint(eff_vram),
|
||||
"outpaint": _select_inpaint(eff_vram), # shares inpaint pipeline
|
||||
"upscale": _select_upscale(eff_vram),
|
||||
}
|
||||
|
||||
|
||||
def _select_txt2img(eff: float) -> ModelSpec:
|
||||
# FLUX.1-schnell (Apache 2.0, 4-step distilled)
|
||||
if eff >= 20.0:
|
||||
return ModelSpec("black-forest-labs/FLUX.1-schnell", "flux", "none", 1024, 20.0)
|
||||
if eff >= 10.0:
|
||||
return ModelSpec("black-forest-labs/FLUX.1-schnell", "flux", "model_cpu_offload", 1024, 20.0)
|
||||
# SDXL base
|
||||
if eff >= 7.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-xl-base-1.0", "sdxl", "none", 1024, 6.5)
|
||||
if eff >= 5.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-xl-base-1.0", "sdxl", "attention_slicing", 1024, 6.5)
|
||||
if eff >= 4.0:
|
||||
return ModelSpec("stabilityai/stable-diffusion-xl-base-1.0", "sdxl", "model_cpu_offload", 1024, 6.5)
|
||||
# SD 2.x
|
||||
if eff >= 3.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-2-1", "sd2x", "none", 768, 3.5)
|
||||
if eff >= 2.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-2-1-base", "sd2x", "attention_slicing", 512, 3.2)
|
||||
# SD 1.5
|
||||
if eff >= 1.7:
|
||||
return ModelSpec("stable-diffusion-v1-5/stable-diffusion-v1-5", "sd15", "attention_slicing", 512, 1.7)
|
||||
return ModelSpec("stable-diffusion-v1-5/stable-diffusion-v1-5", "sd15", "sequential_cpu_offload", 512, 1.7)
|
||||
|
||||
|
||||
def _select_img2img(eff: float) -> ModelSpec:
|
||||
# img2img uses the same model family as txt2img
|
||||
s = _select_txt2img(eff)
|
||||
# FLUX img2img uses a different pipeline class but same model weights
|
||||
return s
|
||||
|
||||
|
||||
def _select_inpaint(eff: float) -> ModelSpec:
|
||||
# No FLUX inpaint pipeline available yet — SDXL is the ceiling
|
||||
if eff >= 7.5:
|
||||
return ModelSpec("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", "sdxl", "none", 1024, 6.5)
|
||||
if eff >= 5.5:
|
||||
return ModelSpec("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", "sdxl", "attention_slicing", 1024, 6.5)
|
||||
if eff >= 4.0:
|
||||
return ModelSpec("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", "sdxl", "model_cpu_offload", 1024, 6.5)
|
||||
if eff >= 3.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-2-inpainting", "sd2x", "none", 512, 3.5)
|
||||
if eff >= 2.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-2-inpainting", "sd2x", "attention_slicing", 512, 3.5)
|
||||
if eff >= 1.7:
|
||||
return ModelSpec("runwayml/stable-diffusion-inpainting", "sd15", "attention_slicing", 512, 1.7)
|
||||
return ModelSpec("runwayml/stable-diffusion-inpainting", "sd15", "sequential_cpu_offload", 512, 1.7)
|
||||
|
||||
|
||||
def _select_upscale(eff: float) -> Optional[ModelSpec]:
|
||||
# SD x4 upscaler — needs ~2 GB fp16 PLUS headroom for the loaded inpaint/txt2img model.
|
||||
# Only enable if eff_vram suggests room for it as a secondary pipeline.
|
||||
if eff >= 6.0:
|
||||
return ModelSpec("stabilityai/stable-diffusion-x4-upscaler", "sd2x", "attention_slicing", 512, 2.0)
|
||||
return None # fall through to Real-ESRGAN
|
||||
|
||||
|
||||
# ── Tier label (display only) ─────────────────────────────────────────────────
|
||||
|
||||
def _tier_label(eff_vram: float) -> str:
|
||||
if eff_vram >= 20: return "flux_full"
|
||||
if eff_vram >= 10: return "flux_offload"
|
||||
if eff_vram >= 7.5: return "sdxl"
|
||||
if eff_vram >= 5.5: return "sdxl_low"
|
||||
if eff_vram >= 4.0: return "sdxl_offload"
|
||||
if eff_vram >= 3.5: return "sd2x"
|
||||
if eff_vram >= 2.5: return "sd2x_low"
|
||||
if eff_vram >= 1.7: return "sd15"
|
||||
return "minimal"
|
||||
|
||||
|
||||
def _caps(tier: str) -> list[str]:
|
||||
base = ["txt2img", "inpaint", "img2img", "outpaint"]
|
||||
if tier in ("flux_full", "flux_offload", "sdxl", "sdxl_low", "sdxl_offload"):
|
||||
return base + ["upscale_diffusion"]
|
||||
return base
|
||||
|
||||
|
||||
# ── Warnings ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def _build_warnings(
|
||||
tier: str, vram_total: float, vram_free: float,
|
||||
cc: str, major: int, minor: int,
|
||||
fp16: bool, bf16: bool, fp8: bool, xf: bool,
|
||||
) -> list[str]:
|
||||
w = []
|
||||
|
||||
if major < 5:
|
||||
w.append(
|
||||
f"GPU compute capability {cc} is not supported by PyTorch 2.x. "
|
||||
"Upgrade to a Kepler/Maxwell-era or newer GPU (CC ≥ 5.0)."
|
||||
)
|
||||
elif major < 6:
|
||||
w.append(
|
||||
f"GPU is Maxwell-era (CC {cc}). fp32 mode — models need 2× VRAM. "
|
||||
"A Pascal GTX 1000-series or newer card enables fp16."
|
||||
)
|
||||
elif not bf16 and tier in ("flux_full", "flux_offload"):
|
||||
w.append(
|
||||
f"GPU CC {cc}: FLUX runs in fp16 (bf16 needs CC ≥ 8.0). "
|
||||
"Results are still good but Ampere/Ada GPUs are faster here."
|
||||
)
|
||||
|
||||
if fp8 and tier in ("flux_full", "flux_offload"):
|
||||
w.append(
|
||||
"FP8 native support detected (Ada Lovelace / Hopper). "
|
||||
"Set HF_MODEL_TXT2IMG=flux-community/flux.1-schnell-fp8 for ~40% VRAM reduction."
|
||||
)
|
||||
|
||||
if tier == "minimal":
|
||||
w.append(
|
||||
f"Very low effective VRAM ({vram_free:.1f} GB free). "
|
||||
"Sequential CPU offload will be used — expect 10–30 min per image."
|
||||
)
|
||||
elif tier == "sdxl_offload":
|
||||
w.append(
|
||||
f"Limited VRAM ({vram_free:.1f} GB free). "
|
||||
"Using SDXL with model_cpu_offload — better quality than SD 2.x, ~30% slower. "
|
||||
"Install xformers or upgrade to ≥5.5 GB effective VRAM for full-speed SDXL."
|
||||
)
|
||||
elif tier in ("sd15", "sd2x_low"):
|
||||
w.append(
|
||||
f"Limited VRAM ({vram_free:.1f} GB free). "
|
||||
"Using SD 1.5/2.x. Upgrade to ≥5.5 GB free for SDXL quality."
|
||||
)
|
||||
|
||||
if xf:
|
||||
w.append(
|
||||
"xformers detected — attention VRAM reduced ~20-30%. "
|
||||
"You may be able to run a higher-tier model than listed."
|
||||
)
|
||||
else:
|
||||
if tier in ("sdxl_low", "sdxl_offload", "sd2x"):
|
||||
w.append(
|
||||
"xformers not installed. Install it (pip install xformers) to reduce "
|
||||
"VRAM usage ~20-30% and potentially unlock the next model tier."
|
||||
)
|
||||
|
||||
return w
|
||||
|
||||
|
||||
# ── Helpers ───────────────────────────────────────────────────────────────────
|
||||
|
||||
def _xformers_available() -> bool:
|
||||
try:
|
||||
import xformers # noqa: F401
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
def _apple_usable_gb() -> float:
|
||||
"""Estimate GPU-usable unified memory (≈ half of total RAM)."""
|
||||
try:
|
||||
r = subprocess.run(
|
||||
["sysctl", "-n", "hw.memsize"], capture_output=True, text=True, timeout=5
|
||||
)
|
||||
if r.returncode == 0:
|
||||
return int(r.stdout.strip()) / (1024 ** 3) / 2
|
||||
except Exception:
|
||||
pass
|
||||
return 8.0
|
||||
|
||||
|
||||
def infer_spec_from_model_id(model_id: str) -> ModelSpec:
|
||||
"""
|
||||
When the user supplies HF_MODEL_* overrides, infer the pipeline family
|
||||
from naming conventions so the correct diffusers class is chosen.
|
||||
"""
|
||||
mid = model_id.lower()
|
||||
if "flux" in mid:
|
||||
return ModelSpec(model_id, "flux", "model_cpu_offload", 1024, 20.0)
|
||||
if "xl" in mid or "sdxl" in mid:
|
||||
return ModelSpec(model_id, "sdxl", "attention_slicing", 1024, 6.5)
|
||||
if any(x in mid for x in ["sd-2", "sd2", "stable-diffusion-2", "-2-", "-2inpaint"]):
|
||||
res = 512 if "base" in mid else 768
|
||||
return ModelSpec(model_id, "sd2x", "attention_slicing", res, 3.5)
|
||||
return ModelSpec(model_id, "sd15", "attention_slicing", 512, 1.7)
|
||||
|
||||
|
||||
# ── Singleton ─────────────────────────────────────────────────────────────────
|
||||
|
||||
_cached: Optional[GpuCapabilities] = None
|
||||
|
||||
|
||||
def get_cached_gpu_info() -> GpuCapabilities:
|
||||
global _cached
|
||||
if _cached is None:
|
||||
_cached = detect_gpu()
|
||||
return _cached
|
||||
|
||||
|
||||
# Alias kept for any callers still using the old name
|
||||
def get_model_ids(tier: str) -> dict:
|
||||
"""Compatibility shim — returns model_id strings keyed by operation."""
|
||||
info = get_cached_gpu_info()
|
||||
return {
|
||||
op: (spec.model_id if spec else None)
|
||||
for op, spec in info.recommended.items()
|
||||
}
|
||||
@@ -0,0 +1,584 @@
|
||||
"""
|
||||
Local GPU diffusion provider — HuggingFace Diffusers backend.
|
||||
|
||||
Implements RemoteAIProvider so all existing routes work unchanged.
|
||||
Pipelines are lazy-loaded, cached in an LRU store, and memory-optimised
|
||||
per the ModelSpec chosen by gpu_detect.
|
||||
|
||||
Supported model families:
|
||||
flux → FluxPipeline / FluxImg2ImgPipeline (FLUX.1-schnell)
|
||||
sdxl → StableDiffusionXL*Pipeline (SDXL base + SDXL Inpaint)
|
||||
sd2x → StableDiffusion2*Pipeline (SD 2.x)
|
||||
sd15 → StableDiffusionPipeline (SD 1.5)
|
||||
|
||||
Requires: diffusers>=0.28.0,<0.29.0, transformers, accelerate, safetensors
|
||||
(all in requirements.gpu.txt — pinned <0.29.0 for PyTorch 2.1.x compatibility)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from collections import OrderedDict
|
||||
from io import BytesIO
|
||||
from typing import Optional
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from app.services.gpu_detect import (
|
||||
GpuCapabilities,
|
||||
ModelSpec,
|
||||
get_cached_gpu_info,
|
||||
infer_spec_from_model_id,
|
||||
)
|
||||
from app.services.remote_provider import RemoteAIProvider
|
||||
|
||||
# ── Model state tracking ──────────────────────────────────────────────────────
|
||||
|
||||
_states: dict[str, dict] = {}
|
||||
_states_lock = threading.Lock()
|
||||
|
||||
|
||||
def _set_state(key: str, **kw):
|
||||
with _states_lock:
|
||||
_states.setdefault(key, {}).update(kw)
|
||||
|
||||
|
||||
def get_all_model_states() -> list[dict]:
|
||||
with _states_lock:
|
||||
return list(_states.values())
|
||||
|
||||
|
||||
def _make_step_cb(pipe_type: str, total_steps: int):
|
||||
"""
|
||||
Returns a diffusers callback_on_step_end that writes per-step progress
|
||||
into _states so the SSE /api/generate/progress endpoint can stream it.
|
||||
Called from a thread executor — _set_state is thread-safe.
|
||||
"""
|
||||
def cb(pipe, step_index: int, timestep, callback_kwargs: dict) -> dict:
|
||||
done = step_index + 1
|
||||
_set_state(pipe_type,
|
||||
state="running",
|
||||
step=done,
|
||||
total_steps=total_steps,
|
||||
progress=round(done / total_steps * 85, 1),
|
||||
message=f"Step {done} / {total_steps}")
|
||||
return callback_kwargs
|
||||
return cb
|
||||
|
||||
|
||||
# ── LRU pipeline cache ────────────────────────────────────────────────────────
|
||||
|
||||
class _PipelineCache:
|
||||
def __init__(self, maxsize: int = 2):
|
||||
self._cache: OrderedDict[str, object] = OrderedDict()
|
||||
self._maxsize = maxsize
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def get(self, key: str):
|
||||
async with self._lock:
|
||||
if key in self._cache:
|
||||
self._cache.move_to_end(key)
|
||||
return self._cache[key]
|
||||
return None
|
||||
|
||||
async def put(self, key: str, pipe: object):
|
||||
async with self._lock:
|
||||
if key in self._cache:
|
||||
self._cache.move_to_end(key)
|
||||
else:
|
||||
if len(self._cache) >= self._maxsize:
|
||||
evicted_key, evicted = self._cache.popitem(last=False)
|
||||
_evict(evicted, evicted_key)
|
||||
self._cache[key] = pipe
|
||||
|
||||
|
||||
def _evict(pipe, key: str):
|
||||
try:
|
||||
import torch
|
||||
pipe.to("cpu")
|
||||
torch.cuda.empty_cache()
|
||||
print(f"[local_gpu] Evicted '{key}' from GPU cache")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ── Pipeline loading helpers ──────────────────────────────────────────────────
|
||||
|
||||
def _apply_hf_token():
|
||||
try:
|
||||
from app.config import settings
|
||||
if settings.hf_token:
|
||||
import huggingface_hub
|
||||
huggingface_hub.login(token=settings.hf_token, add_to_git_credential=False)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _get_spec(pipe_type: str, info: GpuCapabilities) -> ModelSpec:
|
||||
"""Return the ModelSpec for a pipeline type, respecting user overrides."""
|
||||
# Map outpaint to inpaint (same pipeline)
|
||||
op_key = "inpaint" if pipe_type == "outpaint" else pipe_type
|
||||
# img2img uses same family/model as txt2img for FLUX/SDXL
|
||||
if pipe_type == "img2img" and op_key not in info.recommended:
|
||||
op_key = "txt2img"
|
||||
|
||||
# User config override
|
||||
try:
|
||||
from app.config import settings
|
||||
override_map = {
|
||||
"inpaint": settings.hf_model_inpaint,
|
||||
"outpaint": settings.hf_model_inpaint,
|
||||
"txt2img": settings.hf_model_txt2img,
|
||||
"img2img": settings.hf_model_img2img,
|
||||
}
|
||||
override_id = override_map.get(pipe_type, "") or ""
|
||||
if override_id:
|
||||
return infer_spec_from_model_id(override_id)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
spec = info.recommended.get(op_key)
|
||||
if spec is None:
|
||||
raise RuntimeError(
|
||||
f"No model available for '{pipe_type}' at effective VRAM "
|
||||
f"{info.effective_vram_gb:.1f} GB. GPU may not have enough memory."
|
||||
)
|
||||
return spec
|
||||
|
||||
|
||||
def _load_sd_pipeline(pipe_type: str, spec: ModelSpec, info: GpuCapabilities) -> object:
|
||||
"""Load a Stable Diffusion (1.5 / 2.x / XL) pipeline."""
|
||||
import torch
|
||||
from diffusers import (
|
||||
StableDiffusionPipeline,
|
||||
StableDiffusionImg2ImgPipeline,
|
||||
StableDiffusionInpaintPipeline,
|
||||
StableDiffusionUpscalePipeline,
|
||||
StableDiffusionXLPipeline,
|
||||
StableDiffusionXLImg2ImgPipeline,
|
||||
StableDiffusionXLInpaintPipeline,
|
||||
)
|
||||
|
||||
dtype = torch.float16 if info.fp16 else torch.float32
|
||||
is_xl = spec.family == "sdxl"
|
||||
kwargs: dict = {"torch_dtype": dtype}
|
||||
if not is_xl:
|
||||
kwargs["safety_checker"] = None
|
||||
kwargs["requires_safety_checker"] = False
|
||||
|
||||
op_key = "inpaint" if pipe_type in ("inpaint", "outpaint") else pipe_type
|
||||
|
||||
if op_key == "inpaint":
|
||||
cls = StableDiffusionXLInpaintPipeline if is_xl else StableDiffusionInpaintPipeline
|
||||
elif op_key == "txt2img":
|
||||
cls = StableDiffusionXLPipeline if is_xl else StableDiffusionPipeline
|
||||
elif op_key == "img2img":
|
||||
cls = StableDiffusionXLImg2ImgPipeline if is_xl else StableDiffusionImg2ImgPipeline
|
||||
elif op_key == "upscale":
|
||||
cls = StableDiffusionUpscalePipeline
|
||||
else:
|
||||
raise ValueError(f"Unknown SD operation: {op_key}")
|
||||
|
||||
pipe = cls.from_pretrained(spec.model_id, **kwargs)
|
||||
return _apply_mem_opts(pipe, spec, info)
|
||||
|
||||
|
||||
def _load_flux_pipeline(pipe_type: str, spec: ModelSpec, info: GpuCapabilities) -> object:
|
||||
"""Load a FLUX pipeline (txt2img or img2img)."""
|
||||
import torch
|
||||
from diffusers import FluxPipeline, FluxImg2ImgPipeline
|
||||
|
||||
# FLUX works best in bf16 on Ampere+; fp16 on older Turing/Pascal
|
||||
dtype = torch.bfloat16 if info.bf16 else torch.float16
|
||||
|
||||
op_key = "img2img" if pipe_type == "img2img" else "txt2img"
|
||||
cls = FluxImg2ImgPipeline if op_key == "img2img" else FluxPipeline
|
||||
|
||||
pipe = cls.from_pretrained(spec.model_id, torch_dtype=dtype)
|
||||
return _apply_mem_opts(pipe, spec, info)
|
||||
|
||||
|
||||
def _apply_mem_opts(pipe, spec: ModelSpec, info: GpuCapabilities) -> object:
|
||||
"""Apply memory optimisations then move pipeline to device."""
|
||||
device = info.backend
|
||||
opt = spec.memory_opt
|
||||
|
||||
# VAE slicing is always beneficial (reduces VRAM for decoding large images)
|
||||
try:
|
||||
pipe.enable_vae_slicing()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# xformers memory-efficient attention
|
||||
if info.xformers and spec.family != "flux":
|
||||
try:
|
||||
pipe.enable_xformers_memory_efficient_attention()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if opt == "sequential_cpu_offload":
|
||||
# Each layer moved to GPU only during its forward pass — very VRAM-efficient
|
||||
# enable_sequential_cpu_offload() also calls .to(device) internally
|
||||
try:
|
||||
pipe.enable_sequential_cpu_offload()
|
||||
except Exception:
|
||||
pipe.to("cpu")
|
||||
|
||||
elif opt == "model_cpu_offload":
|
||||
# Entire sub-models (text encoder, unet/transformer, VAE) moved between CPU/GPU
|
||||
# Faster than sequential but needs ~3-4 GB free to hold the active module
|
||||
try:
|
||||
pipe.enable_model_cpu_offload()
|
||||
except Exception:
|
||||
pipe.to(device)
|
||||
|
||||
elif opt == "attention_slicing":
|
||||
try:
|
||||
pipe.enable_attention_slicing(1)
|
||||
except Exception:
|
||||
pass
|
||||
pipe.to(device)
|
||||
|
||||
else: # "none"
|
||||
pipe.to(device)
|
||||
|
||||
return pipe
|
||||
|
||||
|
||||
# ── Provider ─────────────────────────────────────────────────────────────────
|
||||
|
||||
class LocalDiffusionProvider(RemoteAIProvider):
|
||||
def __init__(self, max_cached_pipelines: int = 2):
|
||||
self._cache = _PipelineCache(maxsize=max_cached_pipelines)
|
||||
self._load_locks: dict[str, asyncio.Lock] = {}
|
||||
self._meta_lock = asyncio.Lock()
|
||||
|
||||
@property
|
||||
def _info(self) -> GpuCapabilities:
|
||||
return get_cached_gpu_info()
|
||||
|
||||
async def _lock_for(self, key: str) -> asyncio.Lock:
|
||||
async with self._meta_lock:
|
||||
if key not in self._load_locks:
|
||||
self._load_locks[key] = asyncio.Lock()
|
||||
return self._load_locks[key]
|
||||
|
||||
def _load_pipeline_sync(self, pipe_type: str) -> object:
|
||||
info = self._info
|
||||
spec = _get_spec(pipe_type, info)
|
||||
|
||||
_apply_hf_token()
|
||||
_set_state(pipe_type, pipeline=pipe_type, model_id=spec.model_id,
|
||||
family=spec.family, memory_opt=spec.memory_opt,
|
||||
state="downloading", progress=0.0,
|
||||
message=f"Downloading {spec.model_id}…", error="")
|
||||
try:
|
||||
if spec.family == "flux":
|
||||
pipe = _load_flux_pipeline(pipe_type, spec, info)
|
||||
else:
|
||||
pipe = _load_sd_pipeline(pipe_type, spec, info)
|
||||
|
||||
_set_state(pipe_type, state="ready", progress=100.0, message="Ready")
|
||||
return pipe
|
||||
except Exception as exc:
|
||||
_set_state(pipe_type, state="failed", error=str(exc), message="Load failed")
|
||||
raise
|
||||
|
||||
async def _get_pipeline(self, pipe_type: str) -> object:
|
||||
cached = await self._cache.get(pipe_type)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
lock = await self._lock_for(pipe_type)
|
||||
async with lock:
|
||||
cached = await self._cache.get(pipe_type)
|
||||
if cached is not None:
|
||||
return cached
|
||||
loop = asyncio.get_event_loop()
|
||||
pipe = await loop.run_in_executor(None, self._load_pipeline_sync, pipe_type)
|
||||
await self._cache.put(pipe_type, pipe)
|
||||
return pipe
|
||||
|
||||
# ── RemoteAIProvider ──────────────────────────────────────────────────────
|
||||
|
||||
async def inpaint(self, image_bytes: bytes, mask_bytes: bytes, prompt: str, params: dict) -> bytes:
|
||||
pipe = await self._get_pipeline("inpaint")
|
||||
spec = _get_spec("inpaint", self._info)
|
||||
|
||||
img = Image.open(BytesIO(image_bytes)).convert("RGB")
|
||||
mask = Image.open(BytesIO(mask_bytes)).convert("L")
|
||||
orig = img.size
|
||||
img_r, mask_r = _resize_pair(img, mask, spec.native_res)
|
||||
|
||||
steps = int(params.get("steps", 30))
|
||||
cfg = float(params.get("cfg_scale", 7.5))
|
||||
neg = params.get("negative_prompt", "") or None
|
||||
step_cb = _make_step_cb("inpaint", steps)
|
||||
|
||||
_set_state("inpaint", state="running", step=0, total_steps=steps, progress=0, message="Starting…")
|
||||
|
||||
def _run():
|
||||
try:
|
||||
return pipe(
|
||||
prompt=prompt,
|
||||
negative_prompt=neg,
|
||||
image=img_r,
|
||||
mask_image=mask_r,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=cfg,
|
||||
callback_on_step_end=step_cb,
|
||||
callback_on_step_end_tensor_inputs=["latents"],
|
||||
).images[0].resize(orig, Image.LANCZOS)
|
||||
except TypeError:
|
||||
return pipe(
|
||||
prompt=prompt,
|
||||
negative_prompt=neg,
|
||||
image=img_r,
|
||||
mask_image=mask_r,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=cfg,
|
||||
).images[0].resize(orig, Image.LANCZOS)
|
||||
|
||||
result = _to_png(await asyncio.get_event_loop().run_in_executor(None, _run))
|
||||
_set_state("inpaint", state="ready", step=None, total_steps=None, progress=100, message="Ready")
|
||||
return result
|
||||
|
||||
async def txt2img(self, prompt: str, width: int, height: int, params: dict) -> bytes:
|
||||
pipe = await self._get_pipeline("txt2img")
|
||||
spec = _get_spec("txt2img", self._info)
|
||||
|
||||
max_dim = spec.native_res
|
||||
w = min(width, max_dim) // 8 * 8
|
||||
h = min(height, max_dim) // 8 * 8
|
||||
seed = int(params.get("seed", 0))
|
||||
is_flux = spec.family == "flux"
|
||||
steps = 4 if is_flux else int(params.get("steps", 30))
|
||||
step_cb = _make_step_cb("txt2img", steps)
|
||||
|
||||
_set_state("txt2img", state="running", step=0, total_steps=steps, progress=0, message="Starting…")
|
||||
|
||||
def _run():
|
||||
import torch
|
||||
device = self._info.backend
|
||||
gen = torch.Generator(device=device).manual_seed(seed) if seed else None
|
||||
|
||||
try:
|
||||
if is_flux:
|
||||
return pipe(
|
||||
prompt=prompt,
|
||||
width=w, height=h,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=0.0,
|
||||
max_sequence_length=256,
|
||||
generator=gen,
|
||||
callback_on_step_end=step_cb,
|
||||
callback_on_step_end_tensor_inputs=["latents"],
|
||||
).images[0]
|
||||
else:
|
||||
return pipe(
|
||||
prompt=prompt,
|
||||
negative_prompt=params.get("negative_prompt", "") or None,
|
||||
width=w, height=h,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=float(params.get("cfg_scale", 7.5)),
|
||||
generator=gen,
|
||||
callback_on_step_end=step_cb,
|
||||
callback_on_step_end_tensor_inputs=["latents"],
|
||||
).images[0]
|
||||
except TypeError:
|
||||
# Older diffusers without callback_on_step_end
|
||||
if is_flux:
|
||||
return pipe(
|
||||
prompt=prompt, width=w, height=h,
|
||||
num_inference_steps=steps, guidance_scale=0.0,
|
||||
max_sequence_length=256, generator=gen,
|
||||
).images[0]
|
||||
else:
|
||||
return pipe(
|
||||
prompt=prompt,
|
||||
negative_prompt=params.get("negative_prompt", "") or None,
|
||||
width=w, height=h,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=float(params.get("cfg_scale", 7.5)),
|
||||
generator=gen,
|
||||
).images[0]
|
||||
|
||||
result = _to_png(await asyncio.get_event_loop().run_in_executor(None, _run))
|
||||
_set_state("txt2img", state="ready", step=None, total_steps=None, progress=100, message="Ready")
|
||||
return result
|
||||
|
||||
async def img2img(self, image_bytes: bytes, prompt: str, strength: float, params: dict) -> bytes:
|
||||
pipe = await self._get_pipeline("img2img")
|
||||
spec = _get_spec("img2img", self._info)
|
||||
|
||||
img = Image.open(BytesIO(image_bytes)).convert("RGB")
|
||||
orig = img.size
|
||||
img_r = _resize_square(img, spec.native_res)
|
||||
is_flux = spec.family == "flux"
|
||||
steps = 4 if is_flux else int(params.get("steps", 30))
|
||||
step_cb = _make_step_cb("img2img", steps)
|
||||
|
||||
_set_state("img2img", state="running", step=0, total_steps=steps, progress=0, message="Starting…")
|
||||
|
||||
def _run():
|
||||
try:
|
||||
if is_flux:
|
||||
result = pipe(
|
||||
prompt=prompt, image=img_r, strength=strength,
|
||||
num_inference_steps=steps, guidance_scale=0.0,
|
||||
callback_on_step_end=step_cb,
|
||||
callback_on_step_end_tensor_inputs=["latents"],
|
||||
).images[0]
|
||||
else:
|
||||
result = pipe(
|
||||
prompt=prompt,
|
||||
negative_prompt=params.get("negative_prompt", "") or None,
|
||||
image=img_r, strength=strength,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=float(params.get("cfg_scale", 7.5)),
|
||||
callback_on_step_end=step_cb,
|
||||
callback_on_step_end_tensor_inputs=["latents"],
|
||||
).images[0]
|
||||
except TypeError:
|
||||
if is_flux:
|
||||
result = pipe(
|
||||
prompt=prompt, image=img_r, strength=strength,
|
||||
num_inference_steps=steps, guidance_scale=0.0,
|
||||
).images[0]
|
||||
else:
|
||||
result = pipe(
|
||||
prompt=prompt,
|
||||
negative_prompt=params.get("negative_prompt", "") or None,
|
||||
image=img_r, strength=strength,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=float(params.get("cfg_scale", 7.5)),
|
||||
).images[0]
|
||||
return result.resize(orig, Image.LANCZOS)
|
||||
|
||||
result = _to_png(await asyncio.get_event_loop().run_in_executor(None, _run))
|
||||
_set_state("img2img", state="ready", step=None, total_steps=None, progress=100, message="Ready")
|
||||
return result
|
||||
|
||||
async def outpaint(self, image_bytes: bytes, direction: str, size: int, prompt: str) -> bytes:
|
||||
from PIL import ImageDraw
|
||||
|
||||
img = Image.open(BytesIO(image_bytes)).convert("RGB")
|
||||
w, h = img.size
|
||||
|
||||
positions = {
|
||||
"right": ((w + size, h), (0, 0), (w, 0, w + size, h)),
|
||||
"left": ((w + size, h), (size, 0), (0, 0, size, h)),
|
||||
"bottom": ((w, h + size), (0, 0), (0, h, w, h + size)),
|
||||
"top": ((w, h + size), (0, size), (0, 0, w, size)),
|
||||
}
|
||||
new_size, paste_at, mask_box = positions[direction]
|
||||
|
||||
expanded = Image.new("RGB", new_size, (127, 127, 127))
|
||||
expanded.paste(img, paste_at)
|
||||
mask = Image.new("L", new_size, 0)
|
||||
ImageDraw.Draw(mask).rectangle(mask_box, fill=255)
|
||||
|
||||
fill_prompt = prompt or "seamless natural continuation of the scene"
|
||||
return await self.inpaint(_to_png(expanded), _to_png(mask), fill_prompt, {})
|
||||
|
||||
async def health(self) -> bool:
|
||||
return True
|
||||
|
||||
def capabilities(self) -> list[str]:
|
||||
return self._info.capabilities
|
||||
|
||||
|
||||
# ── Image utilities ───────────────────────────────────────────────────────────
|
||||
|
||||
def _resize_pair(img: Image.Image, mask: Image.Image, target: int):
|
||||
w, h = img.size
|
||||
scale = target / max(w, h)
|
||||
nw = max(8, int(w * scale) // 8 * 8)
|
||||
nh = max(8, int(h * scale) // 8 * 8)
|
||||
return img.resize((nw, nh), Image.LANCZOS), mask.resize((nw, nh), Image.NEAREST)
|
||||
|
||||
|
||||
def _resize_square(img: Image.Image, target: int) -> Image.Image:
|
||||
w, h = img.size
|
||||
scale = target / max(w, h)
|
||||
nw = max(8, int(w * scale) // 8 * 8)
|
||||
nh = max(8, int(h * scale) // 8 * 8)
|
||||
return img.resize((nw, nh), Image.LANCZOS)
|
||||
|
||||
|
||||
def _to_png(img: Image.Image) -> bytes:
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
# ── Singleton ─────────────────────────────────────────────────────────────────
|
||||
|
||||
_provider: Optional[LocalDiffusionProvider] = None
|
||||
|
||||
|
||||
def get_local_diffusion_provider(max_pipelines: int = 2) -> LocalDiffusionProvider:
|
||||
global _provider
|
||||
if _provider is None:
|
||||
_provider = LocalDiffusionProvider(max_cached_pipelines=max_pipelines)
|
||||
return _provider
|
||||
|
||||
|
||||
async def prefetch_model_files() -> None:
|
||||
"""
|
||||
Download model weight files to HuggingFace disk cache without loading into GPU.
|
||||
Called at container startup so the first request loads from disk (fast).
|
||||
"""
|
||||
from app.services.gpu_detect import get_cached_gpu_info
|
||||
try:
|
||||
from huggingface_hub import snapshot_download
|
||||
except ImportError:
|
||||
print("[local_gpu] huggingface_hub not installed — skipping model prefetch")
|
||||
return
|
||||
|
||||
info = get_cached_gpu_info()
|
||||
_apply_hf_token()
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
seen: set[str] = set()
|
||||
|
||||
for op, spec in info.recommended.items():
|
||||
if spec is None or spec.model_id in seen:
|
||||
continue
|
||||
seen.add(spec.model_id)
|
||||
|
||||
# Apply user override if set
|
||||
try:
|
||||
from app.config import settings
|
||||
override_map = {
|
||||
"inpaint": settings.hf_model_inpaint,
|
||||
"txt2img": settings.hf_model_txt2img,
|
||||
"img2img": settings.hf_model_img2img,
|
||||
}
|
||||
override = override_map.get(op, "") or ""
|
||||
if override and override not in seen:
|
||||
seen.add(override)
|
||||
spec = infer_spec_from_model_id(override)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
_set_state(op, pipeline=op, model_id=spec.model_id, family=spec.family,
|
||||
memory_opt=spec.memory_opt, state="downloading", progress=0.0,
|
||||
message=f"Downloading {spec.model_id}…", error="")
|
||||
print(f"[local_gpu] Prefetching: {spec.model_id}")
|
||||
|
||||
def _dl(model_id=spec.model_id):
|
||||
snapshot_download(
|
||||
repo_id=model_id,
|
||||
ignore_patterns=["*.msgpack", "flax_*", "tf_*", "rust_model*"],
|
||||
)
|
||||
|
||||
try:
|
||||
await loop.run_in_executor(None, _dl)
|
||||
_set_state(op, state="cached", progress=100.0,
|
||||
message="Files cached — loads into GPU on first request")
|
||||
print(f"[local_gpu] ✓ Cached: {spec.model_id}")
|
||||
except Exception as exc:
|
||||
_set_state(op, state="download_failed", error=str(exc),
|
||||
message="Download failed — will retry on first request")
|
||||
print(f"[local_gpu] Prefetch failed for {spec.model_id}: {exc}")
|
||||
@@ -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,206 @@
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
from PIL import Image
|
||||
from datetime import datetime
|
||||
|
||||
from app.models.patch import Patch
|
||||
from app.config import settings
|
||||
|
||||
|
||||
class PatchLibraryService:
|
||||
"""Service for managing the patch library"""
|
||||
|
||||
def __init__(self, data_dir: str = None):
|
||||
self.data_dir = data_dir or settings.data_dir
|
||||
self.patch_library_dir = Path(self.data_dir) / "patch_library"
|
||||
self.patch_library_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def get_patch_path(self, patch_id: int) -> Path:
|
||||
"""Get path to patch file"""
|
||||
return self.patch_library_dir / f"{patch_id}.png"
|
||||
|
||||
def get_thumbnail_path(self, patch_id: int) -> Path:
|
||||
"""Get path to patch thumbnail"""
|
||||
return self.patch_library_dir / f"{patch_id}_thumb.png"
|
||||
|
||||
def create_thumbnail(self, image_path: Path, thumbnail_path: Path, size: tuple = (200, 200)):
|
||||
"""Create a thumbnail from an image"""
|
||||
img = Image.open(image_path)
|
||||
img.thumbnail(size, Image.Resampling.LANCZOS)
|
||||
img.save(thumbnail_path, 'PNG')
|
||||
|
||||
def save_patch_from_file(
|
||||
self,
|
||||
patch_id: int,
|
||||
image_path: str,
|
||||
create_thumb: bool = True
|
||||
) -> str:
|
||||
"""
|
||||
Save a patch from an existing file
|
||||
|
||||
Args:
|
||||
patch_id: Patch ID
|
||||
image_path: Source image path
|
||||
create_thumb: Whether to create thumbnail
|
||||
|
||||
Returns:
|
||||
Relative path to saved patch
|
||||
"""
|
||||
patch_path = self.get_patch_path(patch_id)
|
||||
shutil.copy(image_path, patch_path)
|
||||
|
||||
if create_thumb:
|
||||
thumbnail_path = self.get_thumbnail_path(patch_id)
|
||||
self.create_thumbnail(patch_path, thumbnail_path)
|
||||
|
||||
return str(patch_path.relative_to(self.data_dir))
|
||||
|
||||
def save_patch_from_bytes(
|
||||
self,
|
||||
patch_id: int,
|
||||
image_bytes: bytes,
|
||||
create_thumb: bool = True
|
||||
) -> str:
|
||||
"""
|
||||
Save a patch from bytes
|
||||
|
||||
Args:
|
||||
patch_id: Patch ID
|
||||
image_bytes: Image data as bytes
|
||||
create_thumb: Whether to create thumbnail
|
||||
|
||||
Returns:
|
||||
Relative path to saved patch
|
||||
"""
|
||||
patch_path = self.get_patch_path(patch_id)
|
||||
|
||||
# Save image
|
||||
with open(patch_path, 'wb') as f:
|
||||
f.write(image_bytes)
|
||||
|
||||
if create_thumb:
|
||||
thumbnail_path = self.get_thumbnail_path(patch_id)
|
||||
self.create_thumbnail(patch_path, thumbnail_path)
|
||||
|
||||
return str(patch_path.relative_to(self.data_dir))
|
||||
|
||||
def save_ai_generated_patch(
|
||||
self,
|
||||
patch_id: int,
|
||||
edit_dir: Path
|
||||
) -> str:
|
||||
"""
|
||||
Save an AI-generated patch from an edit
|
||||
|
||||
Args:
|
||||
patch_id: Patch ID
|
||||
edit_dir: Path to edit history directory
|
||||
|
||||
Returns:
|
||||
Relative path to saved patch
|
||||
"""
|
||||
# Use the AI-generated output (patch_out.png)
|
||||
source_path = edit_dir / "patch_out.png"
|
||||
return self.save_patch_from_file(patch_id, str(source_path))
|
||||
|
||||
def save_manual_patch(
|
||||
self,
|
||||
patch_id: int,
|
||||
project_id: int,
|
||||
bbox: dict
|
||||
) -> str:
|
||||
"""
|
||||
Save a manually selected patch from current project image
|
||||
|
||||
Args:
|
||||
patch_id: Patch ID
|
||||
project_id: Project ID
|
||||
bbox: Bounding box {x, y, width, height}
|
||||
|
||||
Returns:
|
||||
Relative path to saved patch
|
||||
"""
|
||||
from app.services.edit_service import EditService
|
||||
from app.utils.image_processing import crop_patch
|
||||
|
||||
edit_service = EditService(self.data_dir)
|
||||
current_image_path = edit_service.get_current_image_path(project_id)
|
||||
|
||||
# Load and crop current image
|
||||
img = Image.open(current_image_path)
|
||||
patch = crop_patch(img, bbox)
|
||||
|
||||
# Save patch
|
||||
patch_path = self.get_patch_path(patch_id)
|
||||
patch.save(patch_path, 'PNG')
|
||||
|
||||
# Create thumbnail
|
||||
thumbnail_path = self.get_thumbnail_path(patch_id)
|
||||
self.create_thumbnail(patch_path, thumbnail_path)
|
||||
|
||||
return str(patch_path.relative_to(self.data_dir))
|
||||
|
||||
def apply_patch_to_image(
|
||||
self,
|
||||
patch_id: int,
|
||||
target_image: Image.Image,
|
||||
bbox: dict,
|
||||
feather_px: int = 5
|
||||
) -> Image.Image:
|
||||
"""
|
||||
Apply a saved patch to a target image
|
||||
|
||||
Args:
|
||||
patch_id: Patch ID to apply
|
||||
target_image: Target image to apply patch to
|
||||
bbox: Where to place the patch {x, y, width, height}
|
||||
feather_px: Feather radius for blending
|
||||
|
||||
Returns:
|
||||
Image with patch applied
|
||||
"""
|
||||
from app.utils.image_processing import insert_patch, create_feathered_mask
|
||||
from PIL import ImageOps
|
||||
|
||||
# Load patch
|
||||
patch_path = self.get_patch_path(patch_id)
|
||||
patch = Image.open(patch_path).convert('RGBA')
|
||||
|
||||
# Resize patch to match bbox if needed
|
||||
if patch.size != (bbox['width'], bbox['height']):
|
||||
patch = patch.resize((bbox['width'], bbox['height']), Image.Resampling.LANCZOS)
|
||||
|
||||
# Create a soft-edged mask for the patch
|
||||
mask = Image.new('L', patch.size, 255)
|
||||
if feather_px > 0:
|
||||
mask = create_feathered_mask(mask, feather_px)
|
||||
|
||||
# Apply mask to patch
|
||||
patch.putalpha(mask)
|
||||
|
||||
# Insert patch into target image
|
||||
result = insert_patch(target_image, patch, bbox)
|
||||
|
||||
return result
|
||||
|
||||
def delete_patch(self, patch_id: int):
|
||||
"""Delete a patch and its thumbnail"""
|
||||
patch_path = self.get_patch_path(patch_id)
|
||||
thumbnail_path = self.get_thumbnail_path(patch_id)
|
||||
|
||||
if patch_path.exists():
|
||||
patch_path.unlink()
|
||||
|
||||
if thumbnail_path.exists():
|
||||
thumbnail_path.unlink()
|
||||
|
||||
def get_patch_size(self, patch_id: int) -> tuple:
|
||||
"""Get patch dimensions"""
|
||||
patch_path = self.get_patch_path(patch_id)
|
||||
if not patch_path.exists():
|
||||
return (0, 0)
|
||||
|
||||
img = Image.open(patch_path)
|
||||
return img.size
|
||||
@@ -0,0 +1,474 @@
|
||||
"""
|
||||
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 _build_provider(name: str) -> Optional[RemoteAIProvider]:
|
||||
"""Instantiate a named provider from current settings."""
|
||||
from app.config import settings
|
||||
|
||||
name = (name or "").lower().strip()
|
||||
|
||||
if name == "openai":
|
||||
if not settings.openai_api_key:
|
||||
return None
|
||||
return OpenAIRemoteProvider(settings.openai_api_key, settings.openai_model)
|
||||
|
||||
if name == "invokeai":
|
||||
if not settings.invokeai_url:
|
||||
return None
|
||||
return InvokeAIProvider(settings.invokeai_url, settings.invokeai_default_model)
|
||||
|
||||
if name == "comfyui":
|
||||
if not settings.comfyui_url:
|
||||
return None
|
||||
return ComfyUIProvider(settings.comfyui_url, settings.comfyui_default_model)
|
||||
|
||||
if name == "local_gpu":
|
||||
try:
|
||||
from app.services.local_diffusion import get_local_diffusion_provider
|
||||
return get_local_diffusion_provider(max_pipelines=settings.local_gpu_max_pipelines)
|
||||
except (ImportError, AttributeError) as exc:
|
||||
print(f"[local_gpu] Cannot load diffusion provider: {exc}")
|
||||
return None
|
||||
|
||||
return None
|
||||
|
||||
|
||||
# Map operation names to the settings field that holds the override
|
||||
_OP_FIELD = {
|
||||
"inpaint": "ai_provider_inpaint",
|
||||
"txt2img": "ai_provider_txt2img",
|
||||
"img2img": "ai_provider_img2img",
|
||||
"outpaint": "ai_provider_outpaint",
|
||||
}
|
||||
|
||||
|
||||
def get_remote_provider(operation: Optional[str] = None) -> Optional[RemoteAIProvider]:
|
||||
"""
|
||||
Return the provider for a given operation.
|
||||
|
||||
Resolution order:
|
||||
1. Per-operation override (AI_PROVIDER_INPAINT, AI_PROVIDER_TXT2IMG, etc.)
|
||||
2. Global default (AI_PROVIDER)
|
||||
3. None (local-only mode)
|
||||
|
||||
Example .env for mixed setup:
|
||||
AI_PROVIDER=invokeai # default for inpaint/img2img/outpaint
|
||||
AI_PROVIDER_TXT2IMG=openai # use OpenAI only for text-to-image
|
||||
"""
|
||||
from app.config import settings
|
||||
|
||||
if operation and operation in _OP_FIELD:
|
||||
override = getattr(settings, _OP_FIELD[operation], "")
|
||||
if override:
|
||||
provider = _build_provider(override)
|
||||
if provider is not None:
|
||||
return provider
|
||||
# override configured but not usable (missing key/url) — fall through to default
|
||||
|
||||
return _build_provider(settings.ai_provider)
|
||||
@@ -0,0 +1,184 @@
|
||||
"""
|
||||
SAM (Segment Anything Model) service.
|
||||
|
||||
Auto-downloads the ViT-B checkpoint (~375 MB) on first use.
|
||||
Caches the loaded model in memory; re-uses predictor across calls.
|
||||
|
||||
Prediction API:
|
||||
predict_points(image_bytes, points, labels) -> mask_bytes (PNG, white=selected)
|
||||
points: list of (x, y) in original image pixels
|
||||
labels: list of 1 (include) or 0 (exclude), same length as points
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import io
|
||||
import os
|
||||
import urllib.request
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
# ── Model download ────────────────────────────────────────────────────────────
|
||||
|
||||
SAM_DIR = Path("/app/data/models/sam")
|
||||
SAM_FILENAME = "sam_vit_b_01ec64.pth"
|
||||
SAM_URL = f"https://dl.fbaipublicfiles.com/segment_anything/{SAM_FILENAME}"
|
||||
SAM_PATH = SAM_DIR / SAM_FILENAME
|
||||
|
||||
|
||||
class SamInstallState(str, Enum):
|
||||
idle = "idle"
|
||||
downloading = "downloading"
|
||||
done = "done"
|
||||
failed = "failed"
|
||||
|
||||
|
||||
@dataclass
|
||||
class SamInstallStatus:
|
||||
state: SamInstallState = SamInstallState.idle
|
||||
progress: int = 0
|
||||
message: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
_install_status = SamInstallStatus()
|
||||
_install_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def get_install_status() -> dict:
|
||||
s = _install_status
|
||||
return {"state": s.state.value, "progress": s.progress,
|
||||
"message": s.message, "error": s.error}
|
||||
|
||||
|
||||
def sam_model_available() -> bool:
|
||||
return SAM_PATH.exists() and SAM_PATH.stat().st_size > 100_000_000
|
||||
|
||||
|
||||
async def ensure_sam_installed() -> bool:
|
||||
"""Download SAM ViT-B checkpoint if not present. Returns True on success."""
|
||||
global _install_status
|
||||
|
||||
if sam_model_available():
|
||||
_install_status = SamInstallStatus(state=SamInstallState.done, progress=100,
|
||||
message="SAM model ready.")
|
||||
return True
|
||||
|
||||
async with _install_lock:
|
||||
if sam_model_available():
|
||||
_install_status = SamInstallStatus(state=SamInstallState.done, progress=100,
|
||||
message="SAM model ready.")
|
||||
return True
|
||||
|
||||
if _install_status.state == SamInstallState.downloading:
|
||||
return False
|
||||
|
||||
try:
|
||||
SAM_DIR.mkdir(parents=True, exist_ok=True)
|
||||
_install_status = SamInstallStatus(
|
||||
state=SamInstallState.downloading, progress=0,
|
||||
message="Downloading SAM ViT-B model (~375 MB)…",
|
||||
)
|
||||
|
||||
def _download():
|
||||
def _progress(count, block, total):
|
||||
if total > 0:
|
||||
_install_status.progress = min(99, int(count * block * 99 / total))
|
||||
tmp = SAM_PATH.with_suffix(".tmp")
|
||||
urllib.request.urlretrieve(SAM_URL, tmp, _progress)
|
||||
tmp.rename(SAM_PATH)
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, _download)
|
||||
|
||||
_install_status = SamInstallStatus(state=SamInstallState.done, progress=100,
|
||||
message="SAM model ready.")
|
||||
return True
|
||||
|
||||
except Exception as exc:
|
||||
_install_status = SamInstallStatus(
|
||||
state=SamInstallState.failed, error=str(exc),
|
||||
message="SAM download failed.",
|
||||
)
|
||||
print(f"[sam] Download failed: {exc}")
|
||||
return False
|
||||
|
||||
|
||||
# ── Model cache ───────────────────────────────────────────────────────────────
|
||||
|
||||
_predictor = None
|
||||
_predictor_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def _load_predictor():
|
||||
"""Load SAM model and return a SamPredictor. Called in thread pool."""
|
||||
global _predictor
|
||||
if _predictor is not None:
|
||||
return _predictor
|
||||
|
||||
import torch
|
||||
from segment_anything import sam_model_registry, SamPredictor
|
||||
|
||||
if torch.cuda.is_available():
|
||||
device = "cuda"
|
||||
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
device = "mps"
|
||||
else:
|
||||
device = "cpu"
|
||||
|
||||
print(f"[sam] Loading SAM ViT-B on {device}…")
|
||||
sam = sam_model_registry["vit_b"](checkpoint=str(SAM_PATH))
|
||||
sam.to(device)
|
||||
_predictor = SamPredictor(sam)
|
||||
print("[sam] Model loaded.")
|
||||
return _predictor
|
||||
|
||||
|
||||
# ── Prediction ────────────────────────────────────────────────────────────────
|
||||
|
||||
def _predict_sync(image_bytes: bytes,
|
||||
points: list[tuple[int, int]],
|
||||
labels: list[int]) -> bytes:
|
||||
"""
|
||||
Run SAM prediction synchronously (call via run_in_executor).
|
||||
Returns PNG bytes: white = selected, black = background.
|
||||
"""
|
||||
predictor = _load_predictor()
|
||||
|
||||
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
||||
img_array = np.array(image)
|
||||
|
||||
predictor.set_image(img_array)
|
||||
|
||||
pt_array = np.array(points, dtype=np.float32) # [[x, y], ...]
|
||||
lbl_array = np.array(labels, dtype=np.int32) # [1=fg, 0=bg, ...]
|
||||
|
||||
masks, scores, _ = predictor.predict(
|
||||
point_coords=pt_array,
|
||||
point_labels=lbl_array,
|
||||
multimask_output=True,
|
||||
)
|
||||
|
||||
# Pick the highest-confidence mask
|
||||
best = masks[int(np.argmax(scores))] # bool array H×W
|
||||
|
||||
mask_img = Image.fromarray((best * 255).astype(np.uint8), mode="L")
|
||||
buf = io.BytesIO()
|
||||
mask_img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
async def predict_points(image_bytes: bytes,
|
||||
points: list[tuple[int, int]],
|
||||
labels: list[int]) -> bytes:
|
||||
"""Async wrapper for SAM point prediction."""
|
||||
if not sam_model_available():
|
||||
ok = await ensure_sam_installed()
|
||||
if not ok:
|
||||
raise RuntimeError("SAM model not available.")
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, _predict_sync, image_bytes, points, labels)
|
||||
@@ -0,0 +1,494 @@
|
||||
"""
|
||||
Upscale service — auto-detects best available method and runs it.
|
||||
Auto-installs Real-ESRGAN NCNN Vulkan binary when Vulkan GPU is available.
|
||||
Skips NCNN on headless/CPU-only machines and uses PyTorch CPU or Lanczos instead.
|
||||
|
||||
Priority (auto mode):
|
||||
1. Real-ESRGAN PyTorch + CUDA GPU — fastest, best quality
|
||||
2. Real-ESRGAN PyTorch + Apple MPS — fast on Apple Silicon
|
||||
3. Real-ESRGAN NCNN Vulkan binary — fast on any Vulkan GPU
|
||||
4. Real-ESRGAN PyTorch CPU — AI quality, slow (~1-3 min)
|
||||
5. Lanczos — always available, instant
|
||||
|
||||
Capability probe is run once at first call and cached.
|
||||
NCNN binary is auto-downloaded only when Vulkan is detected.
|
||||
Set REALESRGAN_NCNN=force env var to override the Vulkan check.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import shutil
|
||||
import stat
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import urllib.request
|
||||
import zipfile
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from PIL import Image
|
||||
|
||||
# ── NCNN auto-install ─────────────────────────────────────────────────────────
|
||||
|
||||
NCNN_DEST_DIR = Path("/app/data/models/realesrgan")
|
||||
NCNN_VERSION = "v0.2.5.0"
|
||||
NCNN_BASE_URL = f"https://github.com/xinntao/Real-ESRGAN/releases/download/{NCNN_VERSION}"
|
||||
|
||||
_PLATFORM_ZIP = {
|
||||
"linux": f"realesrgan-ncnn-vulkan-{NCNN_VERSION}-ubuntu.zip",
|
||||
"darwin": f"realesrgan-ncnn-vulkan-{NCNN_VERSION}-macos.zip",
|
||||
"win32": f"realesrgan-ncnn-vulkan-{NCNN_VERSION}-windows.zip",
|
||||
"windows": f"realesrgan-ncnn-vulkan-{NCNN_VERSION}-windows.zip",
|
||||
}
|
||||
|
||||
|
||||
class InstallState(str, Enum):
|
||||
idle = "idle"
|
||||
skipped = "skipped" # headless / no Vulkan
|
||||
downloading = "downloading"
|
||||
extracting = "extracting"
|
||||
verifying = "verifying"
|
||||
done = "done"
|
||||
failed = "failed"
|
||||
|
||||
|
||||
@dataclass
|
||||
class InstallStatus:
|
||||
state: InstallState = InstallState.idle
|
||||
progress: int = 0 # 0-100
|
||||
message: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
_install_status = InstallStatus()
|
||||
_install_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def get_install_status() -> dict:
|
||||
s = _install_status
|
||||
return {
|
||||
"state": s.state.value,
|
||||
"progress": s.progress,
|
||||
"message": s.message,
|
||||
"error": s.error,
|
||||
}
|
||||
|
||||
|
||||
def _ncnn_binary_name() -> str:
|
||||
return "realesrgan-ncnn-vulkan.exe" if "win" in sys.platform.lower() else "realesrgan-ncnn-vulkan"
|
||||
|
||||
|
||||
def _vulkan_available() -> bool:
|
||||
"""
|
||||
Check whether a Vulkan-capable GPU is accessible.
|
||||
Returns True if confident a GPU with Vulkan exists; False on headless/CPU-only.
|
||||
Set REALESRGAN_NCNN=force to bypass this check.
|
||||
"""
|
||||
if os.environ.get("REALESRGAN_NCNN", "").lower() == "force":
|
||||
return True
|
||||
|
||||
plat = sys.platform.lower()
|
||||
|
||||
if plat == "linux":
|
||||
# DRI render nodes exist when a GPU is present and drivers loaded
|
||||
dri = Path("/dev/dri")
|
||||
if dri.exists() and list(dri.glob("renderD*")):
|
||||
return True
|
||||
# Fallback: vulkaninfo (not always installed)
|
||||
if shutil.which("vulkaninfo"):
|
||||
r = subprocess.run(["vulkaninfo", "--summary"],
|
||||
capture_output=True, timeout=5)
|
||||
if r.returncode == 0 and b"GPU" in r.stdout:
|
||||
return True
|
||||
return False
|
||||
|
||||
if plat == "darwin":
|
||||
# macOS with Metal/MPS — Vulkan via MoltenVK always present on Apple Silicon/modern Intel
|
||||
return True
|
||||
|
||||
if "win" in plat:
|
||||
# Windows always has a display adapter; assume Vulkan available
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _test_ncnn_binary(binary_path: Path) -> bool:
|
||||
"""Run binary with --help to confirm it actually works (Vulkan loads ok)."""
|
||||
try:
|
||||
r = subprocess.run(
|
||||
[str(binary_path), "--help"],
|
||||
capture_output=True, timeout=15,
|
||||
)
|
||||
# NCNN binary exits 255 for --help but prints usage; that's fine.
|
||||
# A Vulkan init failure produces "no vulkan device" on stderr.
|
||||
stderr = r.stderr.decode(errors="replace").lower()
|
||||
if "no vulkan" in stderr or "failed to create" in stderr:
|
||||
return False
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
async def ensure_ncnn_installed() -> Optional[Path]:
|
||||
"""
|
||||
Check for Vulkan, then download+install the NCNN binary if needed.
|
||||
Skips silently on headless/CPU-only machines.
|
||||
Returns binary Path on success, None otherwise.
|
||||
"""
|
||||
global _install_status
|
||||
|
||||
binary_path = NCNN_DEST_DIR / _ncnn_binary_name()
|
||||
|
||||
# Already installed — quick verify it still works
|
||||
if binary_path.exists() and os.access(binary_path, os.X_OK):
|
||||
loop = asyncio.get_event_loop()
|
||||
ok = await loop.run_in_executor(None, _test_ncnn_binary, binary_path)
|
||||
if ok:
|
||||
_install_status = InstallStatus(state=InstallState.done, progress=100,
|
||||
message="Already installed.")
|
||||
return binary_path
|
||||
else:
|
||||
# Binary exists but Vulkan broken — treat as headless
|
||||
_install_status = InstallStatus(
|
||||
state=InstallState.skipped,
|
||||
message="Vulkan unavailable — skipping NCNN (using PyTorch CPU or Lanczos).",
|
||||
)
|
||||
return None
|
||||
|
||||
async with _install_lock:
|
||||
# Re-check after lock
|
||||
if binary_path.exists() and os.access(binary_path, os.X_OK):
|
||||
_install_status = InstallStatus(state=InstallState.done, progress=100,
|
||||
message="Already installed.")
|
||||
return binary_path
|
||||
|
||||
if _install_status.state in (InstallState.downloading, InstallState.extracting,
|
||||
InstallState.verifying):
|
||||
return None # already running
|
||||
|
||||
# Check Vulkan before downloading anything
|
||||
loop = asyncio.get_event_loop()
|
||||
has_vulkan = await loop.run_in_executor(None, _vulkan_available)
|
||||
if not has_vulkan:
|
||||
_install_status = InstallStatus(
|
||||
state=InstallState.skipped,
|
||||
message="No Vulkan GPU detected — skipping NCNN install. "
|
||||
"AI upscaling via PyTorch CPU or set REALESRGAN_NCNN=force to override.",
|
||||
)
|
||||
print("[upscale] Headless/no-Vulkan detected — skipping NCNN download.")
|
||||
return None
|
||||
|
||||
plat = sys.platform.lower()
|
||||
zip_name = _PLATFORM_ZIP.get(plat)
|
||||
if not zip_name:
|
||||
_install_status = InstallStatus(
|
||||
state=InstallState.failed,
|
||||
error=f"Unsupported platform: {plat}",
|
||||
)
|
||||
return None
|
||||
|
||||
url = f"{NCNN_BASE_URL}/{zip_name}"
|
||||
|
||||
try:
|
||||
NCNN_DEST_DIR.mkdir(parents=True, exist_ok=True)
|
||||
zip_path = NCNN_DEST_DIR / zip_name
|
||||
|
||||
# Download
|
||||
_install_status = InstallStatus(
|
||||
state=InstallState.downloading, progress=0,
|
||||
message=f"Downloading Real-ESRGAN NCNN {NCNN_VERSION}…",
|
||||
)
|
||||
|
||||
def _do_download():
|
||||
def _progress(count, block, total):
|
||||
if total > 0:
|
||||
_install_status.progress = min(85, int(count * block * 85 / total))
|
||||
urllib.request.urlretrieve(url, zip_path, _progress)
|
||||
|
||||
await loop.run_in_executor(None, _do_download)
|
||||
|
||||
# Extract
|
||||
_install_status.state = InstallState.extracting
|
||||
_install_status.progress = 88
|
||||
_install_status.message = "Extracting…"
|
||||
|
||||
def _do_extract():
|
||||
with zipfile.ZipFile(zip_path, "r") as zf:
|
||||
zf.extractall(NCNN_DEST_DIR)
|
||||
found = list(NCNN_DEST_DIR.rglob(_ncnn_binary_name()))
|
||||
if not found:
|
||||
raise FileNotFoundError(f"Binary not found after extract: {_ncnn_binary_name()}")
|
||||
extracted = found[0]
|
||||
if extracted != binary_path:
|
||||
extracted.rename(binary_path)
|
||||
if "win" not in sys.platform.lower():
|
||||
binary_path.chmod(
|
||||
binary_path.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH
|
||||
)
|
||||
zip_path.unlink(missing_ok=True)
|
||||
|
||||
await loop.run_in_executor(None, _do_extract)
|
||||
|
||||
# Verify binary actually works
|
||||
_install_status.state = InstallState.verifying
|
||||
_install_status.progress = 95
|
||||
_install_status.message = "Verifying Vulkan…"
|
||||
|
||||
ok = await loop.run_in_executor(None, _test_ncnn_binary, binary_path)
|
||||
if not ok:
|
||||
binary_path.unlink(missing_ok=True)
|
||||
_install_status = InstallStatus(
|
||||
state=InstallState.skipped,
|
||||
message="Binary installed but Vulkan unavailable at runtime — "
|
||||
"falling back to PyTorch CPU / Lanczos.",
|
||||
)
|
||||
print("[upscale] NCNN binary installed but Vulkan check failed — skipping.")
|
||||
return None
|
||||
|
||||
_install_status = InstallStatus(
|
||||
state=InstallState.done, progress=100,
|
||||
message=f"Real-ESRGAN NCNN installed: {binary_path}",
|
||||
)
|
||||
invalidate_caps_cache()
|
||||
return binary_path
|
||||
|
||||
except Exception as exc:
|
||||
_install_status = InstallStatus(
|
||||
state=InstallState.failed,
|
||||
error=str(exc),
|
||||
message="Installation failed.",
|
||||
)
|
||||
print(f"[upscale] NCNN auto-install failed: {exc}")
|
||||
return None
|
||||
|
||||
|
||||
# ── Capability detection ──────────────────────────────────────────────────────
|
||||
|
||||
_caps: Optional[dict] = None
|
||||
|
||||
|
||||
def probe_upscale_capabilities() -> dict:
|
||||
"""Detect available upscaling methods. Cached after first call."""
|
||||
global _caps
|
||||
if _caps is not None:
|
||||
return _caps
|
||||
|
||||
caps = {
|
||||
"lanczos": True,
|
||||
"realesrgan_pytorch": False,
|
||||
"realesrgan_pytorch_device": None,
|
||||
"realesrgan_ncnn": False,
|
||||
"realesrgan_ncnn_path": None,
|
||||
"recommended": "lanczos",
|
||||
"recommended_label": "Lanczos (no AI upscaler found)",
|
||||
"methods": ["lanczos"],
|
||||
"ncnn_install_status": get_install_status(),
|
||||
}
|
||||
|
||||
# ── PyTorch path ──────────────────────────────────────────────────────────
|
||||
pytorch_device = None
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
pytorch_device = "cuda"
|
||||
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
pytorch_device = "mps"
|
||||
else:
|
||||
pytorch_device = "cpu"
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
if pytorch_device:
|
||||
try:
|
||||
import realesrgan # noqa: F401
|
||||
from basicsr.archs.rrdbnet_arch import RRDBNet # noqa: F401
|
||||
caps["realesrgan_pytorch"] = True
|
||||
caps["realesrgan_pytorch_device"] = pytorch_device
|
||||
caps["methods"].append("realesrgan_pytorch")
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# ── NCNN Vulkan binary ────────────────────────────────────────────────────
|
||||
ncnn_path = _find_ncnn_binary()
|
||||
if ncnn_path:
|
||||
caps["realesrgan_ncnn"] = True
|
||||
caps["realesrgan_ncnn_path"] = str(ncnn_path)
|
||||
caps["methods"].append("realesrgan_ncnn")
|
||||
|
||||
# ── Pick recommended ──────────────────────────────────────────────────────
|
||||
if caps["realesrgan_pytorch"] and pytorch_device in ("cuda", "mps"):
|
||||
device_label = "CUDA GPU" if pytorch_device == "cuda" else "Apple Silicon"
|
||||
caps["recommended"] = "realesrgan_pytorch"
|
||||
caps["recommended_label"] = f"Real-ESRGAN ({device_label})"
|
||||
elif caps["realesrgan_ncnn"]:
|
||||
caps["recommended"] = "realesrgan_ncnn"
|
||||
caps["recommended_label"] = "Real-ESRGAN NCNN (Vulkan)"
|
||||
elif caps["realesrgan_pytorch"] and pytorch_device == "cpu":
|
||||
caps["recommended"] = "realesrgan_pytorch"
|
||||
caps["recommended_label"] = "Real-ESRGAN (CPU — may be slow)"
|
||||
else:
|
||||
install_state = _install_status.state
|
||||
if install_state in (InstallState.downloading, InstallState.extracting, InstallState.verifying):
|
||||
caps["recommended_label"] = "Lanczos (AI upscaler installing…)"
|
||||
elif install_state == InstallState.skipped:
|
||||
caps["recommended_label"] = "Lanczos (headless — no Vulkan GPU)"
|
||||
else:
|
||||
caps["recommended_label"] = "Lanczos (no AI upscaler found)"
|
||||
|
||||
_caps = caps
|
||||
return caps
|
||||
|
||||
|
||||
def _find_ncnn_binary() -> Optional[Path]:
|
||||
found = shutil.which("realesrgan-ncnn-vulkan")
|
||||
if found:
|
||||
return Path(found)
|
||||
candidates = [
|
||||
NCNN_DEST_DIR / _ncnn_binary_name(),
|
||||
Path("/usr/local/bin/realesrgan-ncnn-vulkan"),
|
||||
Path.home() / ".local/bin/realesrgan-ncnn-vulkan",
|
||||
Path(r"C:/realesrgan-ncnn-vulkan/realesrgan-ncnn-vulkan.exe"),
|
||||
Path("/opt/homebrew/bin/realesrgan-ncnn-vulkan"),
|
||||
]
|
||||
for p in candidates:
|
||||
if p.exists() and os.access(p, os.X_OK):
|
||||
return p
|
||||
return None
|
||||
|
||||
|
||||
def invalidate_caps_cache():
|
||||
global _caps
|
||||
_caps = None
|
||||
|
||||
|
||||
# ── Upscale implementations ───────────────────────────────────────────────────
|
||||
|
||||
def _to_png_bytes(img: Image.Image) -> bytes:
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def upscale_lanczos(image: Image.Image, scale: float) -> tuple[bytes, str]:
|
||||
new_w = round(image.width * scale)
|
||||
new_h = round(image.height * scale)
|
||||
result = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
|
||||
return _to_png_bytes(result), "lanczos"
|
||||
|
||||
|
||||
def upscale_realesrgan_pytorch(image: Image.Image, scale: float) -> tuple[bytes, str]:
|
||||
import torch
|
||||
from basicsr.archs.rrdbnet_arch import RRDBNet
|
||||
from realesrgan import RealESRGANer
|
||||
|
||||
caps = probe_upscale_capabilities()
|
||||
device = caps.get("realesrgan_pytorch_device", "cpu")
|
||||
|
||||
model_scale = 2 if scale <= 2.5 else 4
|
||||
model = RRDBNet(
|
||||
num_in_ch=3, num_out_ch=3, num_feat=64,
|
||||
num_block=23, num_grow_ch=32, scale=model_scale
|
||||
)
|
||||
|
||||
model_dir = Path("/app/data/models/realesrgan")
|
||||
model_dir.mkdir(parents=True, exist_ok=True)
|
||||
model_path = model_dir / f"RealESRGAN_x{model_scale}plus.pth"
|
||||
if not model_path.exists():
|
||||
model_path = None
|
||||
|
||||
upsampler = RealESRGANer(
|
||||
scale=model_scale,
|
||||
model_path=str(model_path) if model_path else None,
|
||||
model=model,
|
||||
tile=512,
|
||||
tile_pad=10,
|
||||
pre_pad=0,
|
||||
half=(device == "cuda"),
|
||||
device=torch.device(device),
|
||||
)
|
||||
|
||||
import numpy as np
|
||||
img_bgr = np.array(image)[:, :, ::-1].copy()
|
||||
enhanced, _ = upsampler.enhance(img_bgr, outscale=scale)
|
||||
result = Image.fromarray(enhanced[:, :, ::-1])
|
||||
return _to_png_bytes(result), f"realesrgan_pytorch_{device}"
|
||||
|
||||
|
||||
def upscale_realesrgan_ncnn(image: Image.Image, scale: float) -> tuple[bytes, str]:
|
||||
caps = probe_upscale_capabilities()
|
||||
binary = caps.get("realesrgan_ncnn_path")
|
||||
if not binary:
|
||||
raise RuntimeError("realesrgan-ncnn-vulkan binary not found")
|
||||
|
||||
model_scale = 4 if scale > 2.5 else 2
|
||||
target_w = round(image.width * scale)
|
||||
target_h = round(image.height * scale)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
in_path = Path(tmpdir) / "input.png"
|
||||
out_path = Path(tmpdir) / "output.png"
|
||||
image.save(in_path, format="PNG")
|
||||
|
||||
cmd = [
|
||||
binary,
|
||||
"-i", str(in_path), "-o", str(out_path),
|
||||
"-s", str(model_scale), "-n", f"realesrgan-x{model_scale}plus", "-f", "png",
|
||||
]
|
||||
r = subprocess.run(cmd, capture_output=True, timeout=300)
|
||||
if r.returncode != 0:
|
||||
raise RuntimeError(f"realesrgan-ncnn-vulkan failed: {r.stderr.decode()}")
|
||||
|
||||
result = Image.open(out_path).convert("RGB")
|
||||
if result.width != target_w or result.height != target_h:
|
||||
result = result.resize((target_w, target_h), Image.Resampling.LANCZOS)
|
||||
|
||||
return _to_png_bytes(result), "realesrgan_ncnn"
|
||||
|
||||
|
||||
# ── Public entry point ────────────────────────────────────────────────────────
|
||||
|
||||
def upscale_sync(image: Image.Image, scale: float, method: str = "auto") -> tuple[bytes, str]:
|
||||
"""Upscale synchronously. Returns (png_bytes, method_label)."""
|
||||
caps = probe_upscale_capabilities()
|
||||
|
||||
if method == "auto":
|
||||
method = caps["recommended"]
|
||||
|
||||
if method == "realesrgan_pytorch":
|
||||
if caps["realesrgan_pytorch"]:
|
||||
try:
|
||||
return upscale_realesrgan_pytorch(image, scale)
|
||||
except Exception as e:
|
||||
print(f"Real-ESRGAN PyTorch failed, falling back: {e}")
|
||||
if caps["realesrgan_ncnn"]:
|
||||
try:
|
||||
return upscale_realesrgan_ncnn(image, scale)
|
||||
except Exception as e:
|
||||
print(f"Real-ESRGAN NCNN fallback failed: {e}")
|
||||
return upscale_lanczos(image, scale)
|
||||
|
||||
if method == "realesrgan_ncnn":
|
||||
if caps["realesrgan_ncnn"]:
|
||||
try:
|
||||
return upscale_realesrgan_ncnn(image, scale)
|
||||
except Exception as e:
|
||||
print(f"Real-ESRGAN NCNN failed, falling back: {e}")
|
||||
if caps["realesrgan_pytorch"]:
|
||||
try:
|
||||
return upscale_realesrgan_pytorch(image, scale)
|
||||
except Exception as e:
|
||||
print(f"Real-ESRGAN PyTorch fallback failed: {e}")
|
||||
return upscale_lanczos(image, scale)
|
||||
|
||||
return upscale_lanczos(image, scale)
|
||||
|
||||
|
||||
async def upscale_image(image: Image.Image, scale: float, method: str = "auto") -> tuple[bytes, str]:
|
||||
"""Async wrapper — runs upscale in thread pool."""
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, upscale_sync, image, scale, method)
|
||||
@@ -0,0 +1,228 @@
|
||||
from PIL import Image, ImageFilter, ImageDraw
|
||||
import numpy as np
|
||||
from io import BytesIO
|
||||
from typing import Tuple, Dict
|
||||
import cv2
|
||||
|
||||
|
||||
def bytes_to_image(image_bytes: bytes) -> Image.Image:
|
||||
"""Convert bytes to PIL Image"""
|
||||
return Image.open(BytesIO(image_bytes)).convert('RGBA')
|
||||
|
||||
|
||||
def image_to_bytes(image: Image.Image, format: str = 'PNG') -> bytes:
|
||||
"""Convert PIL Image to bytes"""
|
||||
buffer = BytesIO()
|
||||
image.save(buffer, format=format)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def crop_patch(image: Image.Image, bbox: Dict[str, int]) -> Image.Image:
|
||||
"""
|
||||
Crop a patch from the image using bounding box
|
||||
|
||||
Args:
|
||||
image: PIL Image
|
||||
bbox: Dictionary with x, y, width, height
|
||||
|
||||
Returns:
|
||||
Cropped patch as PIL Image
|
||||
"""
|
||||
x, y, width, height = bbox['x'], bbox['y'], bbox['width'], bbox['height']
|
||||
return image.crop((x, y, x + width, y + height))
|
||||
|
||||
|
||||
def create_feathered_mask(mask: Image.Image, feather_px: int) -> Image.Image:
|
||||
"""
|
||||
Apply feathering (Gaussian blur) to mask edges
|
||||
|
||||
Args:
|
||||
mask: Binary mask image (grayscale)
|
||||
feather_px: Feather radius in pixels
|
||||
|
||||
Returns:
|
||||
Feathered mask
|
||||
"""
|
||||
if feather_px <= 0:
|
||||
return mask
|
||||
|
||||
# Apply Gaussian blur for feathering
|
||||
feathered = mask.filter(ImageFilter.GaussianBlur(radius=feather_px))
|
||||
return feathered
|
||||
|
||||
|
||||
def blend_patch(
|
||||
original_patch: Image.Image,
|
||||
regenerated_patch: Image.Image,
|
||||
mask: Image.Image,
|
||||
feather_px: int = 0,
|
||||
preserve_alpha: bool = True
|
||||
) -> Image.Image:
|
||||
"""
|
||||
Blend regenerated patch with original using mask.
|
||||
Preserves original alpha channel for semi-transparent areas (veils, glass, etc).
|
||||
|
||||
Args:
|
||||
original_patch: Original cropped patch
|
||||
regenerated_patch: AI-regenerated patch
|
||||
mask: Binary mask (same size as patches)
|
||||
feather_px: Feather radius for smooth blending
|
||||
preserve_alpha: If True, preserves original alpha channel
|
||||
|
||||
Returns:
|
||||
Blended patch with preserved transparency
|
||||
"""
|
||||
# Ensure all images are the same size
|
||||
if regenerated_patch.size != original_patch.size:
|
||||
regenerated_patch = regenerated_patch.resize(original_patch.size, Image.Resampling.LANCZOS)
|
||||
|
||||
if mask.size != original_patch.size:
|
||||
mask = mask.resize(original_patch.size, Image.Resampling.LANCZOS)
|
||||
|
||||
# Convert mask to grayscale if needed
|
||||
if mask.mode != 'L':
|
||||
mask = mask.convert('L')
|
||||
|
||||
# Apply feathering to mask
|
||||
feathered_mask = create_feathered_mask(mask, feather_px)
|
||||
|
||||
# Convert images to RGBA, storing original alpha
|
||||
original_rgba = original_patch.convert('RGBA')
|
||||
original_alpha = original_rgba.split()[3] # Store original alpha channel
|
||||
|
||||
regenerated_rgba = regenerated_patch.convert('RGBA')
|
||||
|
||||
# Blend using the feathered mask
|
||||
blended = Image.composite(regenerated_rgba, original_rgba, feathered_mask)
|
||||
|
||||
# Restore original alpha channel to preserve transparency
|
||||
# This keeps semi-transparent areas (veils, glass, smoke) intact
|
||||
if preserve_alpha:
|
||||
r, g, b, _ = blended.split()
|
||||
blended = Image.merge('RGBA', (r, g, b, original_alpha))
|
||||
|
||||
return blended
|
||||
|
||||
|
||||
def insert_patch(
|
||||
full_image: Image.Image,
|
||||
patch: Image.Image,
|
||||
bbox: Dict[str, int]
|
||||
) -> Image.Image:
|
||||
"""
|
||||
Insert a patch back into the full image at the specified bbox
|
||||
|
||||
Args:
|
||||
full_image: Full original image
|
||||
patch: Patch to insert
|
||||
bbox: Bounding box {x, y, width, height}
|
||||
|
||||
Returns:
|
||||
Full image with patch inserted
|
||||
"""
|
||||
result = full_image.copy()
|
||||
x, y = bbox['x'], bbox['y']
|
||||
|
||||
# Ensure patch is the correct size
|
||||
if patch.size != (bbox['width'], bbox['height']):
|
||||
patch = patch.resize((bbox['width'], bbox['height']), Image.Resampling.LANCZOS)
|
||||
|
||||
# Paste the patch
|
||||
result.paste(patch, (x, y), patch if patch.mode == 'RGBA' else None)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def create_mask_from_selection(
|
||||
width: int,
|
||||
height: int,
|
||||
selection_type: str,
|
||||
selection_data: Dict
|
||||
) -> Image.Image:
|
||||
"""
|
||||
Create a binary mask from selection data
|
||||
|
||||
Args:
|
||||
width: Mask width
|
||||
height: Mask height
|
||||
selection_type: "rectangle", "ellipse", or "lasso"
|
||||
selection_data: Selection-specific data
|
||||
|
||||
Returns:
|
||||
Binary mask (white = selected, black = not selected)
|
||||
"""
|
||||
mask = Image.new('L', (width, height), 0)
|
||||
draw = ImageDraw.Draw(mask)
|
||||
|
||||
if selection_type == "rectangle":
|
||||
# Fill entire rectangle
|
||||
draw.rectangle([0, 0, width, height], fill=255)
|
||||
|
||||
elif selection_type == "ellipse":
|
||||
# Fill entire ellipse
|
||||
draw.ellipse([0, 0, width, height], fill=255)
|
||||
|
||||
elif selection_type == "lasso":
|
||||
# Draw polygon from points
|
||||
points = selection_data.get('points', [])
|
||||
if points:
|
||||
# Convert points to relative coordinates within bbox
|
||||
draw.polygon(points, fill=255)
|
||||
|
||||
return mask
|
||||
|
||||
|
||||
def ensure_even_dimensions(image: Image.Image) -> Image.Image:
|
||||
"""
|
||||
Ensure image dimensions are even numbers (required by some AI providers)
|
||||
|
||||
Args:
|
||||
image: PIL Image
|
||||
|
||||
Returns:
|
||||
Image with even dimensions
|
||||
"""
|
||||
width, height = image.size
|
||||
new_width = width if width % 2 == 0 else width + 1
|
||||
new_height = height if height % 2 == 0 else height + 1
|
||||
|
||||
if (new_width, new_height) != (width, height):
|
||||
new_image = Image.new(image.mode, (new_width, new_height), (0, 0, 0, 0))
|
||||
new_image.paste(image, (0, 0))
|
||||
return new_image
|
||||
|
||||
return image
|
||||
|
||||
|
||||
def resize_for_ai(image: Image.Image, max_size: int = 1024) -> Tuple[Image.Image, float]:
|
||||
"""
|
||||
Resize image if needed for AI processing (max dimension)
|
||||
|
||||
Args:
|
||||
image: PIL Image
|
||||
max_size: Maximum dimension size
|
||||
|
||||
Returns:
|
||||
Tuple of (resized image, scale factor)
|
||||
"""
|
||||
width, height = image.size
|
||||
max_dim = max(width, height)
|
||||
|
||||
if max_dim > max_size:
|
||||
scale = max_size / max_dim
|
||||
new_width = int(width * scale)
|
||||
new_height = int(height * scale)
|
||||
resized = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||||
return ensure_even_dimensions(resized), scale
|
||||
|
||||
return ensure_even_dimensions(image), 1.0
|
||||
|
||||
|
||||
def scale_bbox(bbox: Dict[str, int], scale: float) -> Dict[str, int]:
|
||||
"""Scale bounding box coordinates"""
|
||||
return {
|
||||
'x': int(bbox['x'] * scale),
|
||||
'y': int(bbox['y'] * scale),
|
||||
'width': int(bbox['width'] * scale),
|
||||
'height': int(bbox['height'] * scale)
|
||||
}
|
||||
Reference in New Issue
Block a user