Update main.py

main.py

@@ -10,33 +10,30 @@ from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 from fastapi.responses import FileResponse
 
+
+
 app = FastAPI()
 
-# Initialize Gradio Client
-hf_token = os.environ.get('HF_TOKEN')
-client = Client("https://ashrafb-moondream1.hf.space/--replicas/o2fdl/")
-
-# Function to generate captions for uploaded images
-def get_caption(image_path: str, question: str):
-    try:
-        result = client.predict(image_path, question, api_name="/predict")
-        return result
-    except Exception as e:
-        error_msg = "Error occurred while generating caption: {}".format(str(e))
-        raise HTTPException(status_code=500, detail={"error": error_msg})
-
-@app.post("/predict/")
-async def predict(image: UploadFile = File(...), question: str = Form(...)):
-    try:
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_image:
-            shutil.copyfileobj(image.file, temp_image)
-            temp_image_path = temp_image.name
-        result = get_caption(temp_image_path, question)
-        return JSONResponse(content={"result": result})
-    except Exception as e:
-        error_msg = "An unexpected error occurred while processing the request."
-        detail_msg = str(e)
-        raise HTTPException(status_code=500, detail={"error": error_msg, "detail": detail_msg})
+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
+async def save_upload_file(upload_file: UploadFile) -> str:
+    file_path = os.path.join("uploads", upload_file.filename)
+    with open(file_path, "wb") as buffer:
+        buffer.write(await upload_file.read())
+    return file_path
+
+@app.post("/get_caption")
+async def get_caption(image: UploadFile = File(...), context: str = None):
+    # Save the uploaded image to a temporary file
+    file_path = await save_upload_file(image)
+    # Pass the file path to the client for prediction
+    result = client.predict(file_path, context, api_name="/get_caption")
+    return {"caption": result}
+
 app.mount("/", StaticFiles(directory="static", html=True), name="static")
 
 @app.get("/")