api.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import json
from fastapi import FastAPI, Request, HTTPException
import uvicorn

# Determine which handler to use based on environment variable
model_id = os.environ.get("MODEL_ID", "")

if "diffsketcher_edit" in model_id:
    from diffsketcher_edit_handler import Handler
elif "svgdreamer" in model_id:
    from svgdreamer_handler import Handler
elif "diffsketcher" in model_id:
    from diffsketcher_handler import Handler
else:
    raise ValueError(f"Unknown model ID: {model_id}")

app = FastAPI()
handler = Handler()

@app.post("/")
async def process_request(request: Request):
    try:
        json_data = await request.json()
        return handler(json_data)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    return {"status": "ok"}

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