main.py

# main.py
import uvicorn
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from typing import Dict
import os
import time
import shutil

# 只需导入 pipeline
from src.pipeline import main_pipeline

# 不再需要 lifespan，直接创建 app
app = FastAPI(title="Intelligent ID Photo Generator API")

@app.post("/api/v1/idphoto/generate", summary="Generate ID Photo")
async def generate_id_photo(
    user_image: UploadFile = File(..., description="User's portrait photo."),
    template_id: str = Form(..., description="ID of the template to use (e.g., '001').")
) -> Dict:
    
    start_time = time.time()
    
    template_dir = f'assets/templates/{template_id}'
    if not os.path.exists(os.path.join(template_dir, 'template.png')):
        raise HTTPException(status_code=404, detail=f"Template ID '{template_id}' is missing template.png.")

    user_image_path = os.path.join('inputs', user_image.filename)
    with open(user_image_path, "wb") as buffer:
        shutil.copyfileobj(user_image.file, buffer)

    try:
        # --- 调用已修改的 pipeline ---
        results_base64 = main_pipeline(user_image_path, template_id)
        
        end_time = time.time()
        processing_time = round(end_time - start_time, 2)
        
        response_data = {
            "status": "success",
            "processing_time_seconds": processing_time,
            "results": results_base64
        }
        return response_data

    except Exception as e:
        print(f"[!!!] Pipeline Error: {e}")
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if os.path.exists(user_image_path):
            os.remove(user_image_path)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8080)