Mnddoke / main1.py
Ashrafb's picture
Create main1.py
c978c43 verified
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
# Initialize logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = FastAPI()
client = Client("Ashrafb/moondream_captioning")
# Create the "uploads" directory if it doesn't exist
os.makedirs("uploads", exist_ok=True)
# Define a function to save uploaded file to a temporary file
async def save_upload_file(upload_file: UploadFile) -> str:
try:
# Create a temporary directory if it doesn't exist
os.makedirs("temp_uploads", exist_ok=True)
# Create a temporary file path
temp_file_path = os.path.join("temp_uploads", tempfile.NamedTemporaryFile(delete=False).name)
# Save the uploaded file to the temporary file
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=["*"], # Adjust as needed, '*' allows requests from any origin
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.post("/get_caption")
async def get_caption(image: UploadFile = File(...), context: str = Form(...)):
try:
# Save the uploaded image to a temporary file
temp_file_path = await save_upload_file(image)
# Debugging: Print the value of additional_context
logger.info(f"Additional Context: {context}")
# Check if additional context is provided and not None
if context is not None:
context = context.strip()
# Log the parameters being passed to the Gradio client
logger.info(f"Calling client.predict with image={temp_file_path} and context={context}")
# Use handle_file to handle the file upload correctly
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}")
# Serve the app
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)