api.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Dict, Any, Optional, List, Union
import base64
import io
from PIL import Image
import torch
import os
import sys

# Import the handler
from handler import EndpointHandler

# Initialize the app
app = FastAPI()

# Initialize the model
model = EndpointHandler(model_dir="/code")

class TextToImageRequest(BaseModel):
    inputs: Union[str, Dict[str, Any]]
    parameters: Optional[Dict[str, Any]] = None

@app.post("/")
async def text_to_image(request: TextToImageRequest):
    try:
        # Process the request
        result = model(request.dict())
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Add a health check endpoint
@app.get("/health")
async def health():
    return {"status": "ok"}