""" 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="

UI not found

", 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, )