app.py

import os
import shutil
import tempfile
import uuid
import zipfile
from typing import Dict, Any, List

import torch
from fastapi import (
    FastAPI,
    UploadFile,
    File,
    Form,
    HTTPException,
    BackgroundTasks,
)
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from huggingface_hub import hf_hub_download
from PIL import Image
from pydantic import BaseModel

from models import UNet
from test_functions import process_image


# =========================================================
# SETTINGS
# =========================================================

class AppSettings(BaseModel):
    model_repo: str = "Robys01/face-aging"
    model_filename: str = "best_unet_model.pth"
    max_upload_size_mb: int = 20
    allowed_extensions: list[str] = ["jpg", "jpeg", "png", "webp"]


settings = AppSettings()


# =========================================================
# APP
# =========================================================

app = FastAPI(
    title="Face Aging API",
    version="2.0.0",
    description="API-only FastAPI server for face aging image inference."
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["Content-Disposition"]
)


# =========================================================
# MODEL LOAD
# =========================================================

MODEL_DIR = "/tmp/model"
os.makedirs(MODEL_DIR, exist_ok=True)

MODEL_PATH = os.path.join(MODEL_DIR, settings.model_filename)


def download_model() -> None:
    print("Downloading face aging model...")
    hf_hub_download(
        repo_id=settings.model_repo,
        filename=settings.model_filename,
        local_dir=MODEL_DIR,
        cache_dir=os.environ.get("HUGGINGFACE_HUB_CACHE"),
    )


if not os.path.exists(MODEL_PATH):
    download_model()

model = UNet()
model.load_state_dict(
    torch.load(
        MODEL_PATH,
        map_location=torch.device("cpu"),
        weights_only=False,
    )
)
model.eval()
print("Face aging model loaded successfully")


# =========================================================
# UTILITIES
# =========================================================

def validate_image(upload_file: UploadFile) -> None:
    filename = upload_file.filename or ""
    if "." not in filename:
        raise HTTPException(status_code=400, detail="Invalid filename")

    ext = filename.rsplit(".", 1)[-1].lower()
    if ext not in settings.allowed_extensions:
        raise HTTPException(status_code=400, detail="Unsupported image format")


def validate_age(age: int, field_name: str) -> None:
    if age < 0 or age > 100:
        raise HTTPException(
            status_code=400,
            detail=f"{field_name} must be between 0 and 100"
        )


def save_upload_temp(upload_file: UploadFile) -> str:
    suffix = "." + (upload_file.filename or "image.jpg").rsplit(".", 1)[-1]
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)

    with temp_file as buffer:
        shutil.copyfileobj(upload_file.file, buffer)

    return temp_file.name


def save_jpeg_output(image: Image.Image, quality: int = 88) -> str:
    output_filename = f"{uuid.uuid4().hex}.jpg"
    output_path = os.path.join(tempfile.gettempdir(), output_filename)

    if image.mode != "RGB":
        image = image.convert("RGB")

    image.save(
        output_path,
        format="JPEG",
        quality=quality,
        optimize=True
    )

    return output_path


def cleanup_file(path: str) -> None:
    try:
        if path and os.path.exists(path):
            os.remove(path)
    except Exception:
        pass


def parse_target_ages(target_ages: str) -> List[int]:
    """
    Input ví dụ: "40,60"
    Output: [40, 60]
    """
    if not target_ages or not target_ages.strip():
        raise HTTPException(status_code=400, detail="target_ages is required")

    parsed = []
    seen = set()

    for item in target_ages.split(","):
        item = item.strip()
        if not item:
            continue

        try:
            age = int(item)
        except ValueError:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid target age: {item}"
            )

        validate_age(age, "target_age")

        if age not in seen:
            seen.add(age)
            parsed.append(age)

    if not parsed:
        raise HTTPException(status_code=400, detail="No valid target ages provided")

    return parsed


# =========================================================
# ROOT / HEALTH
# =========================================================

@app.get("/")
def root() -> Dict[str, Any]:
    return {
        "status": "ok",
        "message": "Face Aging API is running",
        "endpoints": {
            "health": "GET /health",
            "age_face": "POST /age-face (multipart/form-data: image, source_age, target_age)",
            "age_face_batch": "POST /age-face-batch (multipart/form-data: image, source_age, target_ages='40,60')"
        },
    }


@app.get("/health")
def health() -> Dict[str, Any]:
    return {
        "status": "healthy",
        "model_loaded": True,
        "model_repo": settings.model_repo,
    }


# =========================================================
# SINGLE IMAGE API
# =========================================================

@app.post("/age-face")
async def age_face(
    background_tasks: BackgroundTasks,
    image: UploadFile = File(...),
    source_age: int = Form(...),
    target_age: int = Form(...),
):
    input_path = None
    output_path = None

    try:
        validate_image(image)
        validate_age(source_age, "source_age")
        validate_age(target_age, "target_age")

        input_path = save_upload_temp(image)

        pil_image = Image.open(input_path)
        if pil_image.mode != "RGB":
            pil_image = pil_image.convert("RGB")

        processed_image = process_image(
            model,
            pil_image,
            source_age,
            target_age,
        )

        output_path = save_jpeg_output(processed_image, quality=88)

        background_tasks.add_task(cleanup_file, input_path)
        background_tasks.add_task(cleanup_file, output_path)

        return FileResponse(
            path=output_path,
            media_type="image/jpeg",
            filename="aged_face.jpg",
            headers={
                "Content-Disposition": "inline; filename=aged_face.jpg"
            }
        )

    except HTTPException:
        if input_path:
            cleanup_file(input_path)
        if output_path:
            cleanup_file(output_path)
        raise

    except Exception as e:
        if input_path:
            cleanup_file(input_path)
        if output_path:
            cleanup_file(output_path)

        return JSONResponse(
            status_code=500,
            content={"success": False, "error": str(e)},
        )


# =========================================================
# BATCH API
# =========================================================

@app.post("/age-face-batch")
async def age_face_batch(
    background_tasks: BackgroundTasks,
    image: UploadFile = File(...),
    source_age: int = Form(...),
    target_ages: str = Form(...),   # ví dụ: "40,60"
):
    input_path = None
    zip_path = None

    try:
        validate_image(image)
        validate_age(source_age, "source_age")

        parsed_target_ages = parse_target_ages(target_ages)

        input_path = save_upload_temp(image)

        pil_image = Image.open(input_path)
        if pil_image.mode != "RGB":
            pil_image = pil_image.convert("RGB")

        zip_filename = f"{uuid.uuid4().hex}.zip"
        zip_path = os.path.join(tempfile.gettempdir(), zip_filename)

        with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zipf:
            for target_age in parsed_target_ages:
                processed_image = process_image(
                    model,
                    pil_image,
                    source_age,
                    target_age,
                )

                temp_output_path = save_jpeg_output(processed_image, quality=88)

                zipf.write(
                    temp_output_path,
                    arcname=f"aged_face_{target_age}.jpg"
                )

                cleanup_file(temp_output_path)

        background_tasks.add_task(cleanup_file, input_path)
        background_tasks.add_task(cleanup_file, zip_path)

        return FileResponse(
            path=zip_path,
            media_type="application/zip",
            filename="aged_faces.zip",
            headers={
                "Content-Disposition": "attachment; filename=aged_faces.zip"
            }
        )

    except HTTPException:
        if input_path:
            cleanup_file(input_path)
        if zip_path:
            cleanup_file(zip_path)
        raise

    except Exception as e:
        if input_path:
            cleanup_file(input_path)
        if zip_path:
            cleanup_file(zip_path)

        return JSONResponse(
            status_code=500,
            content={"success": False, "error": str(e)},
        )


# =========================================================
# MAIN
# =========================================================

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "app:app",
        host="0.0.0.0",
        port=int(os.environ.get("PORT", "7860")),
        reload=False,
    )