server/app.py

"""
app.py - FastAPI server for the upgraded local ImageEditter stack.
"""

from __future__ import annotations

import argparse
import base64
import io
import os
import sys
import time
from typing import Optional

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
from PIL import Image, ImageDraw
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, HTMLResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
import uvicorn

from server.providers import EditResult, create_edit_provider
from server.schemas import (
    BackgroundResponse,
    BatchItemResponse,
    BatchResponse,
    CapabilitiesResponse,
    EditResponse,
    GenerateResponse,
    HealthResponse,
    InpaintResponse,
    PresetsResponse,
    StyleTransferResponse,
    UpscaleResponse,
    AdjustResponse,
)


provider = None
device = "cpu"
provider_ready = False


app = FastAPI(title="ImageEditter", description="Hybrid local AI image editor", version="2.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


def _encode_image(image: Image.Image) -> str:
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return base64.b64encode(buffer.getvalue()).decode()


async def _read_upload_image(upload: UploadFile) -> Image.Image:
    data = await upload.read()
    return Image.open(io.BytesIO(data)).convert("RGB")


async def _read_upload_mask(upload: UploadFile) -> Image.Image:
    data = await upload.read()
    return Image.open(io.BytesIO(data)).convert("L")


def _normalize_seed(seed: Optional[int]) -> Optional[int]:
    if seed is None or seed < 0:
        return None
    return seed


def _result_payload(result: EditResult, elapsed: float) -> dict:
    return {
        "success": True,
        "message": f"{result.message} ({elapsed:.1f}s)",
        "image_base64": _encode_image(result.image),
        "elapsed_seconds": elapsed,
        "provider": result.provider,
        "used_fallback": result.used_fallback,
        "steps": result.steps,
        "metadata": result.metadata,
    }


def _provider_info() -> dict:
    if provider is None:
        return {
            "provider": "unavailable",
            "provider_label": "Unavailable",
            "supports_broad_editing": False,
            "supports_sampling_controls": False,
            "supports_generation": False,
            "supports_inpaint": False,
            "supports_batch": False,
            "supports_style_transfer": False,
            "supports_background_ops": False,
            "supports_upscale": False,
            "prompt_hint": "No provider is active.",
            "demo_mode": True,
            "capability_count": 0,
            "preset_count": 0,
        }

    capabilities = provider.capabilities()
    presets = provider.presets()
    supports_diffusion = bool(getattr(provider, "supports_diffusion", False))
    return {
        "provider": provider.provider_id,
        "provider_label": provider.provider_label,
        "supports_broad_editing": provider.supports_broad_editing,
        "supports_sampling_controls": provider.supports_sampling_controls,
        "supports_generation": provider.supports_generation,
        "supports_inpaint": provider.supports_inpaint,
        "supports_batch": provider.supports_batch,
        "supports_style_transfer": provider.supports_style_transfer,
        "supports_background_ops": provider.supports_background_ops,
        "supports_upscale": provider.supports_upscale,
        "prompt_hint": provider.prompt_hint,
        "demo_mode": not supports_diffusion,
        "capability_count": len(capabilities),
        "preset_count": len(presets),
    }


def _ensure_provider():
    if provider is None:
        raise HTTPException(status_code=503, detail="No editing provider is active.")
    return provider


@app.get("/", response_class=HTMLResponse)
async def serve_ui():
    static_dir = os.path.join(os.path.dirname(__file__), "static")
    index_path = os.path.join(static_dir, "index.html")
    if os.path.exists(index_path):
        return FileResponse(index_path, media_type="text/html")
    return HTMLResponse(content="<h1>UI not found</h1>", status_code=404)


@app.get("/health", response_model=HealthResponse)
async def health():
    gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else None
    info = _provider_info()
    return HealthResponse(
        status="ok",
        model_loaded=provider_ready,
        demo_mode=info["demo_mode"],
        device=device,
        gpu_name=gpu_name,
        provider=info["provider"],
        provider_label=info["provider_label"],
        supports_broad_editing=info["supports_broad_editing"],
        supports_sampling_controls=info["supports_sampling_controls"],
        supports_generation=info["supports_generation"],
        supports_inpaint=info["supports_inpaint"],
        supports_batch=info["supports_batch"],
        supports_style_transfer=info["supports_style_transfer"],
        supports_background_ops=info["supports_background_ops"],
        supports_upscale=info["supports_upscale"],
        prompt_hint=info["prompt_hint"],
        capability_count=info["capability_count"],
        preset_count=info["preset_count"],
    )


@app.get("/capabilities", response_model=CapabilitiesResponse)
async def capabilities():
    active_provider = _ensure_provider()
    items = active_provider.capabilities()
    return CapabilitiesResponse(
        provider=active_provider.provider_id,
        provider_label=active_provider.provider_label,
        count=len(items),
        capabilities=items,
    )


@app.get("/presets", response_model=PresetsResponse)
async def presets():
    active_provider = _ensure_provider()
    items = active_provider.presets()
    return PresetsResponse(
        provider=active_provider.provider_id,
        provider_label=active_provider.provider_label,
        count=len(items),
        presets=items,
    )


@app.get("/demo")
async def demo(prompt: str = "make it a rainy night scene with neon reflections"):
    active_provider = _ensure_provider()
    img = Image.new("RGB", (640, 384), color=(40, 58, 84))
    draw = ImageDraw.Draw(img)
    draw.ellipse((210, 52, 430, 244), fill=(222, 180, 150))
    draw.rectangle((260, 210, 380, 380), fill=(48, 120, 88))
    result = active_provider.edit(img, prompt=prompt, num_steps=28)
    buffer = io.BytesIO()
    result.image.save(buffer, format="PNG")
    buffer.seek(0)
    return StreamingResponse(buffer, media_type="image/png")


@app.post("/adjust", response_model=AdjustResponse)
async def adjust_image(
    image: UploadFile = File(...),
    warmth: float = Form(1.0),
    brightness: float = Form(1.0),
    contrast: float = Form(1.0),
    clarity: float = Form(1.0),
    sharpness: float = Form(0.0),
    vignette: float = Form(0.0),
    bloom: float = Form(0.0),
):
    try:
        start = time.time()
        source = await _read_upload_image(image)
        
        from server.cv_engine import CVEngine
        engine = CVEngine()
        
        edited = source.convert("RGB")
        
        # Apply warmth (kelvin shift)
        if warmth != 1.0:
            edited = engine.apply_operation(edited, "white_balance", amount=warmth)
            
        # Apply exposure (brightness)
        if brightness != 1.0:
            from PIL import ImageEnhance
            edited = ImageEnhance.Brightness(edited).enhance(brightness)
            
        # Apply contrast
        if contrast != 1.0:
            from PIL import ImageEnhance
            edited = ImageEnhance.Contrast(edited).enhance(contrast)
            
        # Apply local clarity
        if clarity != 1.0:
            edited = engine.apply_operation(edited, "clarity", amount=clarity)
            
        # Apply sharpening
        if sharpness > 0.0:
            from PIL import ImageFilter
            percent = int(sharpness * 100)
            if percent > 0:
                edited = edited.filter(ImageFilter.UnsharpMask(radius=1.0, percent=percent, threshold=2))
                
        # Apply vignette
        if vignette > 0.0:
            amount = 1.0 - (vignette * 0.6)
            edited = engine.apply_operation(edited, "vignette", amount=amount)
            
        # Apply bloom
        if bloom > 0.0:
            edited = engine.apply_operation(edited, "bloom", amount=bloom)
            
        elapsed = time.time() - start
        return AdjustResponse(
            success=True,
            message=f"Adjusted image properties instantly in {elapsed*1000:.1f}ms",
            image_base64=_encode_image(edited),
            elapsed_seconds=elapsed,
            provider="cv-engine",
            used_fallback=False,
            steps=["adjust"],
        )
    except Exception as exc:
        return AdjustResponse(success=False, message=str(exc))


@app.post("/segment")
async def segment_image(
    image: UploadFile = File(...),
):
    """Magic click-to-select foreground segmentation exactly like Gemini's Magic Editor."""
    active_provider = _ensure_provider()
    try:
        source = await _read_upload_image(image)
        engine = active_provider.engine
        mask_img = engine.segment_foreground(source)
        return {
            "success": True,
            "mask_base64": _encode_image(mask_img),
            "message": "Automatically segmented foreground subject using local GrabCut."
        }
    except Exception as exc:
        return {"success": False, "message": str(exc)}


@app.post("/edit", response_model=EditResponse)
async def edit_image(
    image: UploadFile = File(...),
    prompt: str = Form(...),
    num_steps: int = Form(36),
    text_guidance_scale: float = Form(7.5),
    image_guidance_scale: float = Form(1.5),
    seed: Optional[int] = Form(None),
    mask: Optional[UploadFile] = File(None),
    reference_image: Optional[UploadFile] = File(None),
    background_image: Optional[UploadFile] = File(None),
):
    active_provider = _ensure_provider()

    try:
        start = time.time()
        source = await _read_upload_image(image)
        mask_image = await _read_upload_mask(mask) if mask is not None else None
        reference = await _read_upload_image(reference_image) if reference_image is not None else None
        background = await _read_upload_image(background_image) if background_image is not None else None

        result = active_provider.edit(
            image=source,
            prompt=prompt,
            num_steps=num_steps,
            text_guidance_scale=text_guidance_scale,
            image_guidance_scale=image_guidance_scale,
            seed=_normalize_seed(seed),
            mask=mask_image,
            reference_image=reference,
            background_image=background,
        )
        return EditResponse(**_result_payload(result, time.time() - start))
    except Exception as exc:
        return EditResponse(success=False, message=str(exc))


@app.post("/generate", response_model=GenerateResponse)
async def generate_image(
    prompt: str = Form(...),
    width: int = Form(768),
    height: int = Form(768),
    num_steps: int = Form(40),
    text_guidance_scale: float = Form(7.5),
    image_guidance_scale: float = Form(1.0),
    seed: Optional[int] = Form(None),
):
    active_provider = _ensure_provider()
    try:
        start = time.time()
        result = active_provider.generate(
            prompt=prompt,
            width=width,
            height=height,
            num_steps=num_steps,
            text_guidance_scale=text_guidance_scale,
            image_guidance_scale=image_guidance_scale,
            seed=_normalize_seed(seed),
        )
        return GenerateResponse(**_result_payload(result, time.time() - start))
    except Exception as exc:
        return GenerateResponse(success=False, message=str(exc))


@app.post("/inpaint", response_model=InpaintResponse)
async def inpaint_image(
    image: UploadFile = File(...),
    mask: UploadFile = File(...),
    prompt: str = Form("repair the masked region naturally"),
    num_steps: int = Form(36),
    text_guidance_scale: float = Form(7.5),
    image_guidance_scale: float = Form(1.5),
    seed: Optional[int] = Form(None),
):
    active_provider = _ensure_provider()
    try:
        start = time.time()
        source = await _read_upload_image(image)
        mask_image = await _read_upload_mask(mask)
        result = active_provider.inpaint(
            image=source,
            mask=mask_image,
            prompt=prompt,
            num_steps=num_steps,
            text_guidance_scale=text_guidance_scale,
            image_guidance_scale=image_guidance_scale,
            seed=_normalize_seed(seed),
        )
        return InpaintResponse(**_result_payload(result, time.time() - start))
    except Exception as exc:
        return InpaintResponse(success=False, message=str(exc))


@app.post("/batch", response_model=BatchResponse)
async def batch_edit(
    images: list[UploadFile] = File(...),
    prompt: str = Form(...),
    num_steps: int = Form(36),
    text_guidance_scale: float = Form(7.5),
    image_guidance_scale: float = Form(1.5),
    seed: Optional[int] = Form(None),
):
    active_provider = _ensure_provider()
    try:
        start = time.time()
        pil_images = [await _read_upload_image(item) for item in images]
        results = active_provider.batch_edit(
            images=pil_images,
            prompt=prompt,
            num_steps=num_steps,
            text_guidance_scale=text_guidance_scale,
            image_guidance_scale=image_guidance_scale,
            seed=_normalize_seed(seed),
        )

        items = []
        for idx, result in enumerate(results):
            items.append(
                BatchItemResponse(
                    index=idx,
                    success=True,
                    message=result.message,
                    image_base64=_encode_image(result.image),
                    provider=result.provider,
                    used_fallback=result.used_fallback,
                    steps=result.steps,
                    metadata=result.metadata,
                )
            )

        return BatchResponse(
            success=True,
            message=f"Batch edit complete for {len(items)} image(s).",
            items=items,
            elapsed_seconds=time.time() - start,
            provider=active_provider.provider_id,
        )
    except Exception as exc:
        return BatchResponse(success=False, message=str(exc))


@app.post("/style-transfer", response_model=StyleTransferResponse)
async def style_transfer(
    image: UploadFile = File(...),
    reference_image: UploadFile = File(...),
    prompt: str = Form(""),
    seed: Optional[int] = Form(None),
):
    active_provider = _ensure_provider()
    try:
        start = time.time()
        source = await _read_upload_image(image)
        reference = await _read_upload_image(reference_image)
        result = active_provider.style_transfer(
            image=source,
            reference_image=reference,
            prompt=prompt,
            seed=_normalize_seed(seed),
        )
        return StyleTransferResponse(**_result_payload(result, time.time() - start))
    except Exception as exc:
        return StyleTransferResponse(success=False, message=str(exc))


@app.post("/background", response_model=BackgroundResponse)
async def background_edit(
    image: UploadFile = File(...),
    prompt: str = Form(...),
    background_image: Optional[UploadFile] = File(None),
    seed: Optional[int] = Form(None),
):
    active_provider = _ensure_provider()
    try:
        start = time.time()
        source = await _read_upload_image(image)
        background = await _read_upload_image(background_image) if background_image is not None else None
        result = active_provider.background_edit(
            image=source,
            prompt=prompt,
            background_image=background,
            seed=_normalize_seed(seed),
        )
        return BackgroundResponse(**_result_payload(result, time.time() - start))
    except Exception as exc:
        return BackgroundResponse(success=False, message=str(exc))


@app.post("/upscale", response_model=UpscaleResponse)
async def upscale_image(
    image: UploadFile = File(...),
    scale: float = Form(2.0),
    prompt: str = Form(""),
    seed: Optional[int] = Form(None),
):
    active_provider = _ensure_provider()
    try:
        start = time.time()
        source = await _read_upload_image(image)
        result = active_provider.upscale(
            image=source,
            scale=scale,
            prompt=prompt,
            seed=_normalize_seed(seed),
        )
        return UpscaleResponse(**_result_payload(result, time.time() - start))
    except Exception as exc:
        return UpscaleResponse(success=False, message=str(exc))


static_dir = os.path.join(os.path.dirname(__file__), "static")
if os.path.exists(static_dir):
    app.mount("/static", StaticFiles(directory=static_dir), name="static")


def start_server(
    checkpoint_path: str = None,
    vae_checkpoint_path: str = None,
    host: str = "0.0.0.0",
    port: int = 8000,
    provider_name: str = "auto",
    foundation_backend: str = None,
    foundation_model_id: str = None,
):
    global provider, device, provider_ready

    device = "cuda" if torch.cuda.is_available() else "cpu"

    try:
        provider = create_edit_provider(
            provider_name=provider_name,
            foundation_backend=foundation_backend,
            foundation_model_id=foundation_model_id,
            checkpoint_path=checkpoint_path,
            vae_checkpoint_path=vae_checkpoint_path,
            device=device,
        )
        provider_ready = provider.warmup()
        print(f"Using provider: {provider.provider_label}")
    except Exception as exc:
        if provider_name == "auto":
            print(f"WARNING: Primary provider failed to initialize: {exc}")
            print("Falling back to the local CV engine.")
            provider = create_edit_provider(provider_name="basic", device=device)
            provider_ready = provider.warmup()
        else:
            raise

    uvicorn.run(app, host=host, port=port)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=str, default=None)
    parser.add_argument("--vae-checkpoint", type=str, default=None)
    parser.add_argument("--host", type=str, default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument("--provider", type=str, default="auto", choices=["auto", "basic", "custom", "foundation"])
    parser.add_argument("--foundation-backend", type=str, default=None)
    parser.add_argument("--foundation-model-id", type=str, default=None)
    args = parser.parse_args()

    start_server(
        checkpoint_path=args.checkpoint,
        vae_checkpoint_path=args.vae_checkpoint,
        host=args.host,
        port=args.port,
        provider_name=args.provider,
        foundation_backend=args.foundation_backend,
        foundation_model_id=args.foundation_model_id,
    )