main.py

"""
Face Swap + Gemini API
======================
Merged from python-api/main.py + python-api/python-api/main.py
All Railway deployment fixes applied.

Endpoints
─────────
GET  /            → health check
GET  /health      → health check (alias)
POST /swap        → InsightFace face-swap (lazy-loaded, won't block startup)
POST /gemini/generate → Gemini multimodal generation (legacy SDK)
POST /ai          → Gemini text generation (new SDK, from ai_service)
POST /face        → face detection stub (from face_service)
"""

from __future__ import annotations

import base64
import os
import sys
import gc

# Silence chatty libraries BEFORE imports
os.environ["ORT_LOGGING_LEVEL"] = "3"
os.environ["KMP_WARNINGS"] = "0"
os.environ["INSIGHTFACE_HOME"] = "/app/models"
os.environ["PYTHONUNBUFFERED"] = "1"

import traceback
import logging
import warnings
from pathlib import Path
from typing import Optional

logging.getLogger("onnxruntime").setLevel(logging.ERROR)
logging.getLogger("insightface").setLevel(logging.ERROR)
warnings.filterwarnings("ignore")

# print(">>> NEW VERSION DEPLOYED - Python AI API starting up", flush=True)

import numpy as np
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel

# Service modules (your original files, also fixed)
from ai_service import generate_ai_response
from face_service import process_face

# ── App ───────────────────────────────────────────────────────────────────────

app = FastAPI(title="Face Swap + Gemini API", version="2.0.0")

# Allow requests from any origin (Netlify, localhost, etc.)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global exception handler — always return JSON, never bare HTML 500 pages
@app.exception_handler(Exception)
async def _global_exc_handler(request: Request, exc: Exception) -> JSONResponse:
    # traceback.print_exc(file=sys.stderr)
    return JSONResponse(
        status_code=500,
        content={"error": str(exc) or "Internal server error"},
    )


# ─────────────────────────────────────────────────────────────────────────────
# InsightFace — optional, guarded import
# ─────────────────────────────────────────────────────────────────────────────

MODEL_DIR    = Path(__file__).parent / "models"
SWAPPER_PATH = MODEL_DIR / "inswapper_128.onnx"

# Module-level handles — populated at startup
_face_analyser = None
_face_swapper  = None
_celeb_cache: dict[str, object] = {}

try:
    import cv2                        # noqa: F401 — import check only
    import insightface                # noqa: F401
    _INSIGHTFACE_AVAILABLE = True
    
    from insightface.app import FaceAnalysis
    from insightface.model_zoo import get_model
    
    # print("⏳ Loading FaceAnalysis (buffalo_sc) once globally...", flush=True)
    _face_analyser = FaceAnalysis(
        name="buffalo_sc",
        root=".",
        providers=["CPUExecutionProvider"]
    )
    _face_analyser.prepare(
        ctx_id=-1,
        det_size=(192, 192)
    )
    # print("✅ FaceAnalysis ready", flush=True)

except ImportError:
    _INSIGHTFACE_AVAILABLE = False
    print("⚠️  insightface/cv2 not installed — /swap will return 503", file=sys.stderr)
except Exception as e:
    print(f"❌ Error loading models at startup: {e}", flush=True)


def _load_inswapper():
    """Lazily load inswapper model if not already loaded."""
    global _face_swapper
    if _face_swapper is not None:
        return _face_swapper
    
    from insightface.model_zoo import get_model
    swapper_path = "./models/inswapper_128.onnx"
    if not os.path.exists(swapper_path):
        swapper_path = "/app/models/inswapper_128.onnx"
    if not os.path.exists(swapper_path):
        swapper_path = str(MODEL_DIR / "inswapper_128.onnx")

    # print(f"⏳ Loading inswapper lazily from: {swapper_path}", flush=True)
    _face_swapper = get_model(
        swapper_path,
        download=False,
        providers=["CPUExecutionProvider"]
    )
    
    return _face_swapper


# ── Image helpers (your original code, unchanged) ─────────────────────────────

def b64_to_img(b64: str) -> "np.ndarray":
    import cv2
    if "," in b64:
        b64 = b64.split(",", 1)[1]
    data = base64.b64decode(b64)
    arr  = np.frombuffer(data, np.uint8)
    img  = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError("Could not decode image")
    
    # Optimization: shrink incoming images to prevent OOM
    h, w = img.shape[:2]
    max_size = 720
    if max(h, w) > max_size:
        scale = max_size / max(h, w)
        img = cv2.resize(
            img,
            (int(w * scale), int(h * scale))
        )
        
    return img


def img_to_b64(img: "np.ndarray", quality: int = 85) -> str:
    import cv2
    _, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, quality])
    return base64.b64encode(buf).decode()


def best_face(analyser, img: "np.ndarray"):
    faces = analyser.get(img)
    if not faces:
        return None
    return max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))


# ── Pydantic models ───────────────────────────────────────────────────────────

class SwapRequest(BaseModel):
    source_key: str
    source_b64: Optional[str] = None
    target_b64: str
    quality: int = 55


class AIRequest(BaseModel):
    prompt: str


class GeminiPart(BaseModel):
    text: Optional[str] = None
    inlineData: Optional[dict] = None


class GeminiRequest(BaseModel):
    systemInstruction: str
    parts: list[GeminiPart]


# ─────────────────────────────────────────────────────────────────────────────
# Routes
# ─────────────────────────────────────────────────────────────────────────────

# ── Health ────────────────────────────────────────────────────────────────────

@app.get("/")
@app.get("/health")
def health():
    face_models_in_memory = _face_analyser is not None and _face_swapper is not None
    return {
        "status": "ok",
        # NOTE: models are lazy-loaded on first /swap POST, so this is False
        # until the first swap request.  We separately expose server_ready=True
        # so the frontend can distinguish "server up but models not yet warm"
        # from "server unreachable".
        "models_loaded": face_models_in_memory,
        # Always True while the server is running — frontend uses this for the
        # reachability check instead of models_loaded.
        "server_ready": True,
        "insightface_available": _INSIGHTFACE_AVAILABLE,
        "cached_faces": list(_celeb_cache.keys()),
        "gemini_configured": bool(os.getenv("GEMINI_API_KEY")),
    }


@app.get("/debug")
def debug():
    return {
        "models_folder_exists": os.path.exists("/app/models"),
        "models": os.listdir("/app/models") if os.path.exists("/app/models") else "not found",
        "inswapper_exists": os.path.exists("/app/models/inswapper_128.onnx")
    }


# ── Face swap (your original logic, now loaded globally) ─────────────────────

@app.post("/swap")
async def swap(req: SwapRequest):
    global _face_analyser
    if _face_analyser is None:
        raise HTTPException(503, "FaceAnalysis model is not loaded on this server.")
    
    # Lazy load swapper
    swapper = _load_inswapper()

    # 1. Cache celebrity face if not already done
    if req.source_key not in _celeb_cache:
        if not req.source_b64:
            raise HTTPException(400, "source_b64 required for first call with this key")
        src_img = b64_to_img(req.source_b64)
        face    = best_face(_face_analyser, src_img)
        if face is None:
            raise HTTPException(422, "No face detected in source image")
        _celeb_cache[req.source_key] = face

    source_face = _celeb_cache[req.source_key]

    # 2. Decode target frame
    if not req.target_b64:
        raise HTTPException(400, "target_b64 required")
    target_img   = b64_to_img(req.target_b64)
    target_faces = _face_analyser.get(target_img)

    # 3. No face in target → return original unchanged
    if not target_faces:
        return {"swapped": False, "frame": img_to_b64(target_img, req.quality)}

    # 4. Swap the largest detected face
    target_face = max(
        target_faces,
        key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]),
    )
    result = swapper.get(target_img, target_face, source_face, paste_back=True)
    
    # Free memory after heavy processing
    gc.collect()
    
    return {"swapped": True, "frame": img_to_b64(result, req.quality)}


# ── Gemini multimodal (your original /gemini/generate endpoint, fixed) ────────

@app.post("/gemini/generate")
async def generate_gemini(req: GeminiRequest):
    # Guard: package must be installed
    try:
        import google.generativeai as genai
    except ImportError:
        raise HTTPException(503, "google-generativeai is not installed.")

    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        raise HTTPException(500, "GEMINI_API_KEY is not configured on Railway.")

    genai.configure(api_key=api_key)

    # FIX: gemini-3.1-pro-preview does not exist — use env var or safe default
    model_name = os.getenv("GEMINI_MODEL", "gemini-1.5-pro")
    model = genai.GenerativeModel(model_name)

    try:
        # Build content parts (your original logic, unchanged)
        from google.generativeai import types
        content_parts = []
        for p in req.parts:
            if p.text:
                content_parts.append(types.Part(text=p.text))
            elif p.inlineData:
                m_type = p.inlineData.get("mimeType") or p.inlineData.get("mime_type")
                m_data = p.inlineData.get("data")
                if m_type and m_data:
                    content_parts.append(
                        types.Part(
                            inline_data=types.Blob(
                                mime_type=m_type,
                                data=m_data,
                            )
                        )
                    )

        response = model.generate_content(
            contents=content_parts,
            generation_config=genai.GenerationConfig(temperature=0.1),
            system_instruction=req.systemInstruction,
        )
        return {"text": getattr(response, "text", str(response))}

    except Exception as exc:
        raise HTTPException(500, str(exc)) from exc


# ── Simple AI text endpoint (from your original python-api/main.py) ───────────

@app.post("/ai")
def ai(req: AIRequest):
    """Thin wrapper around ai_service.generate_ai_response."""
    return {"result": generate_ai_response(req.prompt)}


# ── Face stub endpoint (from your original python-api/main.py) ───────────────
@app.post("/face")
def face(data: dict):
    """Thin wrapper around face_service.process_face."""
    return {"result": process_face(data)}


if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port, access_log=False)