server.py

import io
import os
import uuid
import logging
from typing import Optional

from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, Header
from fastapi.responses import FileResponse, JSONResponse
from pydantic import BaseModel
import torch

import numpy as np
from PIL import Image

# Lazy import performed in get_model() to avoid import-time failures on Space


LOGGER = logging.getLogger("hair_server")
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")

EXPECTED_BEARER = "logicgo@123"

# Optional Mongo persistence
from pymongo import MongoClient
MONGO_URI = os.environ.get("MONGO_URI", "")
mongo_client = MongoClient(MONGO_URI) if MONGO_URI else None
mongo_db = mongo_client.get_database("HairSwapDB") if mongo_client else None
uploads_col = mongo_db.get_collection("uploads") if mongo_db else None
results_col = mongo_db.get_collection("results") if mongo_db else None


def verify_bearer(authorization: Optional[str] = Header(None)):
    if not authorization:
        raise HTTPException(status_code=401, detail="Missing Authorization header")
    try:
        scheme, token = authorization.split(" ", 1)
    except ValueError:
        raise HTTPException(status_code=401, detail="Invalid Authorization header format")
    if scheme.lower() != "bearer":
        raise HTTPException(status_code=401, detail="Invalid auth scheme")
    if token != EXPECTED_BEARER:
        raise HTTPException(status_code=401, detail="Invalid token")
    return True


app = FastAPI(title="Hair Swap API", version="1.0.0")


@app.get("/health")
def health():
    return {"status": "healthy"}


@app.get("/")
def root():
    return {"status": "ok"}


class HairSwapRequest(BaseModel):
    source_id: str
    reference_id: str
    converter_scale: float = 1.0
    scale: float = 1.0
    guidance_scale: float = 1.5
    controlnet_conditioning_scale: float = 1.0


# Initialize model lazily on first request
_model = None  # type: ignore[assignment]


def get_model():
    global _model
    if _model is None:
        try:
            LOGGER.info("Loading StableHair model ...")
            device = "cuda" if torch.cuda.is_available() else "cpu"
            dtype = torch.float16 if device == "cuda" else torch.float32
            LOGGER.info(f"Using device: {device}, dtype: {dtype}")
            
            # Import here to defer importing diffusers/transformers until needed
            from infer_full import StableHair  # noqa: WPS433
            _model = StableHair(config="./configs/hair_transfer.yaml", device=device, weight_dtype=dtype)
            LOGGER.info("Model loaded successfully")
        except Exception as e:
            LOGGER.error(f"Failed to load model: {str(e)}")
            raise Exception(f"Model loading failed: {str(e)}")
    return _model


# Use a writable location on Hugging Face Spaces
BASE_DATA_DIR = os.environ.get("DATA_DIR", "/data")
UPLOAD_DIR = os.path.join(BASE_DATA_DIR, "uploads")
RESULTS_DIR = os.path.join(BASE_DATA_DIR, "results")
LOGS_DIR = os.path.join(BASE_DATA_DIR, "logs")
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(RESULTS_DIR, exist_ok=True)
os.makedirs(LOGS_DIR, exist_ok=True)


@app.post("/upload")
async def upload_image(image: UploadFile = File(...), _=Depends(verify_bearer)):
    if not image.filename:
        raise HTTPException(status_code=400, detail="No file name provided")
    contents = await image.read()
    try:
        Image.open(io.BytesIO(contents)).convert("RGB")
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid image file")

    image_id = str(uuid.uuid4())
    ext = os.path.splitext(image.filename)[1] or ".png"
    path = os.path.join(UPLOAD_DIR, image_id + ext)
    with open(path, "wb") as f:
        f.write(contents)
    # Save metadata to Mongo
    if uploads_col:
        try:
            uploads_col.insert_one({"_id": image_id, "filename": os.path.basename(path), "path": path})
        except Exception:
            pass
    return {"id": image_id, "filename": os.path.basename(path)}


@app.post("/get-hairswap")
def get_hairswap(req: HairSwapRequest, _=Depends(verify_bearer)):
    try:
        # Resolve file paths
        def find_file(image_id: str) -> str:
            for name in os.listdir(UPLOAD_DIR):
                if name.startswith(image_id):
                    return os.path.join(UPLOAD_DIR, name)
            raise HTTPException(status_code=404, detail=f"Image id not found: {image_id}")

        source_path = find_file(req.source_id)
        reference_path = find_file(req.reference_id)
        
        LOGGER.info(f"Found source: {source_path}, reference: {reference_path}")

        # Load model with error handling
        try:
            model = get_model()
            LOGGER.info("Model loaded successfully")
        except Exception as e:
            LOGGER.error(f"Model loading failed: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Model loading failed: {str(e)}")

        # Perform hair transfer with error handling
        try:
            LOGGER.info("Starting hair transfer...")
            id_np, out_np, bald_np, ref_np = model.Hair_Transfer(
                source_image=source_path,
                reference_image=reference_path,
                random_seed=-1,
                step=30,
                guidance_scale=req.guidance_scale,
                scale=req.scale,
                controlnet_conditioning_scale=req.controlnet_conditioning_scale,
                size=512,
            )
            LOGGER.info("Hair transfer completed successfully")
        except Exception as e:
            LOGGER.error(f"Hair transfer failed: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Hair transfer failed: {str(e)}")

        # Save result
        try:
            result_id = str(uuid.uuid4())
            out_img = Image.fromarray((out_np * 255.).astype(np.uint8))
            filename = f"{result_id}.png"
            out_path = os.path.join(RESULTS_DIR, filename)
            out_img.save(out_path)
            LOGGER.info(f"Result saved: {out_path}")
            
            if results_col:
                try:
                    results_col.insert_one({
                        "_id": result_id,
                        "filename": filename,
                        "path": out_path,
                        "source_id": req.source_id,
                        "reference_id": req.reference_id,
                    })
                except Exception as e:
                    LOGGER.warning(f"MongoDB save failed: {str(e)}")

            return {"result": filename}
        except Exception as e:
            LOGGER.error(f"Result saving failed: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Result saving failed: {str(e)}")
            
    except HTTPException:
        raise
    except Exception as e:
        LOGGER.error(f"Unexpected error in get_hairswap: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")


@app.get("/download/{filename}")
def download(filename: str, _=Depends(verify_bearer)):
    path = os.path.join(RESULTS_DIR, filename)
    if not os.path.exists(path):
        raise HTTPException(status_code=404, detail="File not found")
    return FileResponse(path, media_type="image/png", filename=filename)


@app.get("/logs")
def logs(_=Depends(verify_bearer)):
    if uploads_col and results_col:
        uploads = list(uploads_col.find({}, {"_id": 1, "filename": 1}).limit(20))
        results = list(results_col.find({}, {"_id": 1, "filename": 1, "source_id": 1, "reference_id": 1}).limit(20))
        return JSONResponse({"uploads": uploads, "results": results})
    return JSONResponse({"logs": ["service running"], "db": "not_configured"})