|
from fastapi import FastAPI, File, UploadFile, Form, HTTPException |
|
from fastapi.responses import JSONResponse |
|
from gradio_client import Client, handle_file |
|
import os |
|
import tempfile |
|
import shutil |
|
from fastapi.middleware.cors import CORSMiddleware |
|
import logging |
|
|
|
|
|
logging.basicConfig(level=logging.INFO) |
|
logger = logging.getLogger(__name__) |
|
|
|
app = FastAPI() |
|
|
|
client = Client("Ashrafb/moondream_captioning") |
|
|
|
|
|
os.makedirs("uploads", exist_ok=True) |
|
|
|
|
|
async def save_upload_file(upload_file: UploadFile) -> str: |
|
try: |
|
|
|
os.makedirs("temp_uploads", exist_ok=True) |
|
|
|
temp_file_path = os.path.join("temp_uploads", tempfile.NamedTemporaryFile(delete=False).name) |
|
|
|
with open(temp_file_path, "wb") as buffer: |
|
shutil.copyfileobj(upload_file.file, buffer) |
|
return temp_file_path |
|
except Exception as e: |
|
logger.error(f"Error saving upload file: {e}") |
|
raise HTTPException(status_code=500, detail=f"Error saving upload file: {e}") |
|
|
|
app.add_middleware( |
|
CORSMiddleware, |
|
allow_origins=["*"], |
|
allow_credentials=True, |
|
allow_methods=["*"], |
|
allow_headers=["*"], |
|
) |
|
|
|
@app.post("/get_caption") |
|
async def get_caption(image: UploadFile = File(...), context: str = Form(...)): |
|
try: |
|
|
|
temp_file_path = await save_upload_file(image) |
|
|
|
|
|
logger.info(f"Additional Context: {context}") |
|
|
|
|
|
if context is not None: |
|
context = context.strip() |
|
|
|
|
|
logger.info(f"Calling client.predict with image={temp_file_path} and context={context}") |
|
|
|
|
|
result = client.predict( |
|
image=handle_file(temp_file_path), |
|
question=context, |
|
api_name="/get_caption" |
|
) |
|
|
|
return {"caption": result} |
|
except Exception as e: |
|
logger.error(f"Error in get_caption: {e}") |
|
raise HTTPException(status_code=500, detail=f"Error in get_caption: {e}") |
|
|
|
|
|
if __name__ == "__main__": |
|
import uvicorn |
|
uvicorn.run(app, host="0.0.0.0", port=7860) |